What is the shape of your states? Depending on the dimensions you used for your network, the shape of the states batch should be torch.Size([batch_size, number of colour channels, width, height]) and each state should be torch.Size([dimension, number of colour channels, width, height])   (assuming you are using images for input) . Where dimension should match the dimensions set when creating the network. Dimension should be 1 to match with the “get current” method
UPDATE*******************
Yes. You could use “state = state.squeeze().unsqueeze(dim=0)” to change the dimension of your state that will agree with the “get current” method. This should be done before pushing the state to memory