If anyone ever comes by I ended up solving the issue by replacing the line
tag_scores, states = prednet(tensor, states)

with
tag_scores, states = prednet(tensor, [ state.detach() for state in states])

worked well ! it is similar solution to the one of @ptrblck just that states is a list of tensor
Thank you for the help!