If you have a prediction function as-
def predict(features):
    with torch.no_grad():
        return model(features)

Then it will always call this context-manager method, which in turn might affect inference performance.
So after training your model, I think it makes more sense to use “torch.set_grad_enabled(False)”, function since it disables any gradient computation globally, so then you can simple create your prediction function as-
def predict(features):
    return model(features)

Am I right @albanD?