Just wanted to add a tip on your model for efficiency. In the particular case of using Relu and maxpool you can actually reorder as doing the maxpool2D on conv2d output and then apply the Relu activation. Reason being that order of operation doesn’t matter in this particular case as they mathematically compute to be equal. So in forward def put order of operations as:
out =F.relu(F.max_pool2d(self.conv1(x), 2))
When done in this order it requires a substantial fewer number of computations and obviously speeds training up as benefit while result is still the same