That’s an error in your code in line 79,80, with input variable. Altought you shouldn’t do a squeeze before go as an input of the model, because model expects a batch dimension
edit: i didn’t notice that it’s an unsqueeze instead of squeeze