I am using the following versions:
python                   3.8.5           h5fd99cc_1
pytorch                  1.8.1           py3.8_cuda11.1_cudnn8_0
torchaudio             0.8.1           py38
torchvision             0.9.1           py38_cu111
For the model, I am using a simple MLP network with 4 hidden layers. I posted the entire source code above.
The code just for the network is:

class Classifier(nn.Module):
def init (self):
super().init ()
self.fc0 = nn.Linear(50176, 784)
self.fc1 = nn.Linear(784, 256)
self.fc2 = nn.Linear(256, 128)
self.fc3 = nn.Linear(128, 64)
self.fc4 = nn.Linear(64, 2)
def forward(self, x):
# make sure input tensor is flattened
print('The shape of X is: ')
print(x.shape)
x = x.view(x.shape[0], -1)
print('The shape of X flattened is: ')
print(x.shape)
x = F.relu(self.fc0(x))
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = F.relu(self.fc3(x))
x = F.log_softmax(self.fc4(x), dim=1)

print('The shape of X is: ')
print(x.shape)

return x