In Training, swapaxis and squeeze is OK for loss.backward()?

autograd
1

In my code, I have many swapaxis and squeeze in the nn class layer.
for example,

    def forward(self, pixel, video_length):
        for avg_pooling, model in enumerate(self.batch):
            pixel = model(pixel)
            if avg_pooling == 1:
                avg_pooling_2d = nn.functional.adaptive_avg_pool2d(pixel, 1)
                avg_pooling_2d = torch.swapaxes(pixel, 2, 0)
                avg_pooling_2d = torch.swapaxes(avg_pooling_2d, 1, 2)
                avg_pooling_2d = torch.squeeze(pixel)
                avg_pooling_2d = torch.unsqueeze(avg_pooling_2d)
                input = self.ReLU(avg_pooling_2d)
                return input

My problem is, when loss.backward is applied, the backpropagation process is automatically recover the swapaxis and squeeze process which is originally located at the def forward() ?

In my opinion… the loss is gradually increasing when the training epoch is increasing… So I doubt that the swapaxis in the training session may broke the dimension order at the backpropagation process. Anyone knows that swapaxis and squeeze is OK for loss.backward in the PyTorch? Please help!!

SIMPLE SUMMURY : Can I use swapaxis in forward()? (no problem in loss.backward() ?)

Thanks.