Hi,
If you try set the grad to false for a layer inside with grad enabled, does this override the with clause?
For example:
with torch.set_grad_enabled(mode = True):
    self.model.lin2.requires_grad = False # Will this work?
   for m in self.model.mlp_f: m.set_grad_enabled = False # Or this?