You need to permute your channels.
Currently you have: (H, W, C, N) but you need it in (N, C, H, W)
If it’s a tensor:
x = x.permute(3, 2, 1, 0)
Should do it. You can put this in your training loop before you feed the variable to your model.