I assume you need the retain_graph=True setting, since you are not detaching the states tensor.
If this is your use case, you would have to lower the batch size to be able to store all computation graphs on the device.
If you don’t need to backpropagate through multiple steps, you might want to detach states via:
tag_scores, states = prednet(tensor, states.detach())