Putting with torch.cuda.amp.autocast(enabled=False): before dec_t = self.dec_t(quant_t) doesn’t change anything, but if I put it in forward pass of the decoder I now get the error RuntimeError: Function 'CudnnConvolutionBackward' returned nan values in its 1th output. and the traceback points to the forward pass. Sorry if I’m doing this wrong, I’m not quite sure where I should disable autocast.