AssertionError: No inf checks were recorded for this optimizer. in custom diffusion model

1

I am trying to modify the UNet2DModel from the diffusers library by creating a new class and inheriting from it. I am trying to attach encoders and decoders onto the diffusion layers. However, I am not quite sure how to change the forward method. This is the code I’ve tried:

class CustomUNet2DModel(UNet2DModel):
    def __init__(self, encoder, decoder, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, x, timesteps, *args, **kwargs):
        # Encode the input
        encoded_x = self.encoder(x)
        
        # Apply the diffusion process using the parent class
        diffusion_output = super().forward(encoded_x, timesteps, *args, **kwargs)
        
        # Decode the output back to the original size
        decoded_output = self.decoder(diffusion_output[0])
        
        return decoded_output, 0

But it’s giving me the error AssertionError: No inf checks were recorded for this optimizer.. I think this is because something about my logic of calculating the forward method is wrong, but I am not quite sure how to fix this.

2

Could you add the missing parts of the code showing how amp is used and how the optimizer is created?

3

It turns out that the problem was because I did not change the model variable in the optimizer definition…
optimizer = torch.optim.AdamW(model.parameters(), lr=TrainingConfig.learning_rate)