Cleaner way to use "with torch.no_grad()" conditioned on an expression
My code looks like:

if no_grad_condition:
  with torch.no_grad():
    out=network(input)
else:
  out=network(input)


Is there a cleaner way to do it, without duplicating the line out=network(input)?

I am looking for something in the spirit of:

  with torch.no_grad(no_grad_condition):
    out=network(input)


OP here: By writing down the question, I understood where to look for the answer. According to pytorch docs, we can use set_grad_enabled:

  with torch.set_grad_enabled(not no_grad_condition):
    out=network(input)

