batch = torch.randn(32, 3, 28, 28)   # N, C, H, W
flat = batch.flatten(1)              # flatten from dim 1 onward
print(flat.shape)                    # torch.Size([32, 2352])