@ptrblck  Yes, your checking code passed (torch summary does not make error on shapes), but my training;
def train(model, device, loader, optimizer):
    model.train()
    total_loss = 0
    correct = 0
    for x, y in loader:
        x = x.permute(0,3,1,2)
        #x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        output = model(x)
        loss = nn.functional.nll_loss(output, y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        y_pred = output.argmax(dim=1, keepdim=True)
        correct += y_pred.eq(y.view_as(y_pred)).sum().item()
    return {
        "loss": total_loss / len(loader.dataset),
        "accuracy": correct / len(loader.dataset),
    }

and
optimizer = torch.optim.Adadelta(model.parameters(), lr=LEARNING_RATE)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.7)
with tqdm(range(1, EPOCH_SIZE)) as progress:
    for epoch in progress:
        result = train(model, device, dataloader_train, optimizer)
        scheduler.step()
        progress.write(str(result))

These generate an error;
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-36-71a7a55dd184> in <module>
      3 with tqdm(range(1, EPOCH_SIZE)) as progress:
      4     for epoch in progress:
----> 5         result = train(model, device, dataloader_train, optimizer)
      6         scheduler.step()
      7         progress.write(str(result))

<ipython-input-35-b7caa5450bda> in train(model, device, loader, optimizer)
      7         #x, y = x.to(device), y.to(device)
      8         optimizer.zero_grad()
----> 9         output = model(x)
     10         loss = nn.functional.nll_loss(output, y)
     11         loss.backward()

~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

<ipython-input-24-893b6a73f9d4> in forward(self, image)
     38         flattened   = torch.flatten(b3_pl, start_dim=1)
     39 
---> 40         b4_fc1      = torch.relu(self.fc1(flattened))
     41         b4_fc2      = torch.relu(self.fc2(b4_fc1))
     42         b4_fc3      = torch.relu(self.fc3(b4_fc2))

~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

~/.local/lib/python3.8/site-packages/torch/nn/modules/linear.py in forward(self, input)
     91 
     92     def forward(self, input: Tensor) -> Tensor:
---> 93         return F.linear(input, self.weight, self.bias)
     94 
     95     def extra_repr(self) -> str:

~/.local/lib/python3.8/site-packages/torch/nn/functional.py in linear(input, weight, bias)
   1688     if input.dim() == 2 and bias is not None:
   1689         # fused op is marginally faster
-> 1690         ret = torch.addmm(bias, input, weight.t())
   1691     else:
   1692         output = input.matmul(weight.t())

RuntimeError: mat1 and mat2 shapes cannot be multiplied (32x19950 and 1587x529)

Maybe I have confusion and misunderstanding on it…