ayaz-amin:

W = self.conv.weight.squeeze()


This will remove all singleton dimensions. As you have 1 channel, you delete all of them from a (1, 1, 1, 1)-shaped-Tensor. You only want to delete the last two.
In general, I would recommend to furnish such code with comments or asserts on the expected shapes at each step. - The shape things are quite messy.
(The code also could use a PyTorch >= 0.4 migration 1, but that’s another story.)
Best regards
Thomas