qiminchen:

But in training, pytorch is in the default channel first so is swapping the channel to last, doing matmul and swapping channel back to first the only solution?


The matrix multiplication doesn’t “check” for the memory format, i.e. if the tensors are in channels-first or channels-last, but just uses both tensors in the provided shapes.
It’s on the user to make sure the result is expected, so I would recommend to create tensors with known values and compare the output to this reference.