Solved by Martin_Barry in post #8 
              
              
                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 tenso…