Solved by soulitzer in post #2 
              
              
                Unfortunately it was removed as part of a refactor to the code. I’m not sure there’s an equivalent you can just import, but since its relatively simple, you might want to just replicate it in your own code. 
This was the original function: 
def zero_gradients(x):
    if isinstance(x, torch.Tensor):
…