Thank you @ptrblck for answering my topic and picking up as well the issue I posted on github ! I put more details regarding amp in the github issue so I will try to answer you more details on this thread.
About the error when calling the stft operation. Since you answered about the difference for CPU and GPU tensors, I manually checked on ipython and indeed calling stft on a FP16 cpu tensor raises an error which is not the case for the same sent on a cuda device.
However, my code runs always on GPU. The line that causes the error in the loss calculation is
mag_grains = torch.norm(torch.stft(mb_grains,n_fft,hop_length=hop_size,win_length=win_size,window=torch.hann_window(win_size).to(device,non_blocking=True),center=False),dim=3)
mb_grains.type() == ‘torch.cuda.HalfTensor’ and device(type=‘cuda’, index=1) so everything should be in GPU memory … I dont understand what is the CPU data causing the error here, do you see it ?
About amp, are you okay with continuing the discussion on the github issue please ?
ps: upper in the code I put torch.set_default_dtype(torch.float16) to make sure that the window is created with the right dtype (as it is not automatically casted by amp in that case if I understood correct)