RuntimeError: Trying to backward through the graph a second time.

[isabel-s-rosa]

Hello!

I am calling a deployed nequip model from C++ using libtorch, but after I run one step of the physics engine and call the model again on the updated input, I see an error that looks like this:

Traceback of TorchScript, original code (most recent call last):
  File "/Users/isabelrosa/Library/Python/3.8/lib/python/site-packages/nequip/nn/_rescale.py", line 124, in forward
    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
        data = self.model(data)
               ~~~~~~~~~~ <--- HERE
        if self.training:
            return data
  File "/Users/isabelrosa/Library/Python/3.8/lib/python/site-packages/nequip/nn/_grad_output.py", line 77, in forward
        data = self.func(data)
        # Get grads
        grads = torch.autograd.grad(
                ~~~~~~~~~~~~~~~~~~~ <--- HERE
            # TODO:
            # This makes sense for scalar batch-level or batch-wise outputs, specifically because d(sum(batches))/d wrt = sum(d batch / d wrt) = d my_batch / d wrt
RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.

Has anyone seen this before or have any idea how to go about fixing it? All I'm doing in my code is calling the forward function on the same model twice with different inputs. Thanks!

[Linux-cpp-lisp]

HI @isabel-s-rosa;

Hm... my first guess:

did you update your input tensor instead of making a new tensor?

[isabel-s-rosa]

My code for running the model is here, it's very messy still, but if you get a chance to look it over and let me know if you see any clear issues, that would be great!

[Linux-cpp-lisp]

Without looking through more closely can already ask: is any of the input to the second step connected in the autograd graph to the outputs of the first? For example, are the new positions a PyTorch operation of old positions and predicted forces?

You can check for this kind of problem by checking the requires_grad in input2 for all the tensors--- it shouldn't be true for any, but some might be true if they accidentally got connected in the autograd graph. To avoid this, use .detach().

[isabel-s-rosa]

That's a good point! I just checked requires_grad on all the tensors that go into the second round of inputs, and it seems like everything is okay. Any other ideas? Thanks for all your help!

[isabel-s-rosa]

Never mind! This was my fault, I was passing in the wrong dictionary, thank you for all your help:)

[Linux-cpp-lisp]

Ah ok, glad to hear you found the issue!