----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 32, 80, 80]             896
              ReLU-2           [-1, 32, 80, 80]               0
         MaxPool2d-3           [-1, 32, 40, 40]               0
            Conv2d-4           [-1, 64, 40, 40]          18,496
              ReLU-5           [-1, 64, 40, 40]               0
         MaxPool2d-6           [-1, 64, 20, 20]               0
           Flatten-7                [-1, 25600]               0
            Linear-8                  [-1, 256]       6,553,600
              ReLU-9                  [-1, 256]               0
          Dropout-10                  [-1, 256]               0
           Linear-11                   [-1, 31]           7,967
================================================================
Total params: 6,580,959
Trainable params: 6,580,959
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.07
Forward/backward pass size (MB): 5.47
Params size (MB): 25.10
Estimated Total Size (MB): 30.65
----------------------------------------------------------------