DSX:

flat = torch.cat([convs[-1]], dim=1).reshape(-1,)


The problem is in this line, you are flatting the vector with the batch dimension.
Try:
cat = torch.cat([convs[-1]], dim=1)
flat = cat.view(cat .size(0), -1)

This should fix that.