Method apply_torch_model

model_helper.py:57–65

(self)

        }

    def apply_torch_model(self):
        def disable_checkpoint(self):
            if getattr(self, "use_checkpoint", False) == True:
                self.use_checkpoint = False
            if getattr(self, "checkpoint", False) == True:
                self.checkpoint = False

        self.unet.apply(disable_checkpoint)
        self.set_unet("None")

    def set_unet(self, ckpt: str):
        # TODO test if using this with TRT works