Method forward

diff2flow/trainer_module.py:194–195

(self, x1: Tensor, x0: Tensor, **kwargs)

        return out

    def forward(self, x1: Tensor, x0: Tensor, **kwargs):
        return self.model.training_losses(x1=x1, x0=x0, **kwargs).mean()

    def add_lora_to_unet(self):
        unet = self.model