Method zero_grad

src/diffusion/fp16_util.py:173–174

(self)

            self.model.convert_to_fp16()

    def zero_grad(self):
        zero_grad(self.model_params)

    def backward(self, loss: th.Tensor):
        if self.use_fp16: