I have few spectrograms which look like this:
x = torch.load('input')
fig, axs = plt.subplots(2, 2, figsize=(24,8))
axs = axs.flat
for idx, i in enumerate(x.cpu().detach().numpy()):
    i = scaler.fit_transform(i)
    axs[idx].imshow(i, origin='lower', cmap='jet')
    axs[idx].axis('off')
    
fig.tight_layout()
fig.suptitle('Correct image', size=30)


image1720×558 318 KB

However, my neural network requires some axes swapping along the last two dimension. I realized that swapping the images distort the images,
x = x.transpose(-1,-2)

fig, axs = plt.subplots(2, 2, figsize=(24,8))
axs = axs.flat
for idx, i in enumerate(x.cpu().detach().numpy()):
    i = scaler.fit_transform(i)
    axs[idx].imshow(i.swapaxes(0,1), origin='lower', cmap='jet')
    axs[idx].axis('off')
    
fig.tight_layout()
fig.suptitle('Wrong image', size=30)


image1720×558 552 KB

What am I doing wrong here?
You can download my images below
https://sutdapac-my.sharepoint.com/:u:/g/personal/kinwai_cheuk_mymail_sutd_edu_sg/Eb711s15rkFDhcTTUh4vRXkB55rZ3WCY_OS2f0rDsbK8qg?e=SR94eu