如何让训练好的神经网络在识别后同时返回运行时间

我用的mobilenetv2,如何在返回结果时同时返回计算所需的时间呢？代码怎么写？python3.6，环境是pytorch

from time import time
start = time()
history=d1_model_ya.fit(X_train, Y_train, epochs=10,batch_size=32, validation_split=0.2)
end = time()
print("CPU_time =" + str(end-start))