Update model.py

Saquib764 · web-flow · commit 17a1b834bb87 · 2025-08-27T13:28:21.000+05:30
diff --git a/src/train/model.py b/src/train/model.py
@@ -35,7 +35,8 @@ def __init__(
             FluxOminiKontextPipeline.from_pretrained(flux_pipe_id).to(dtype=dtype).to(device)
         )
         self.transformer = self.flux_pipe.transformer
-        self.transformer.gradient_checkpointing = gradient_checkpointing
+        if gradient_checkpointing:
+            self.transformer.enable_gradient_checkpointing()
         self.transformer.train()
 
         # Freeze the Flux pipeline components

Original file line number	Diff line number	Diff line change
`@@ -35,7 +35,8 @@ def __init__(`
`35`	`35`	`FluxOminiKontextPipeline.from_pretrained(flux_pipe_id).to(dtype=dtype).to(device)`
`36`	`36`	`)`
`37`	`37`	`self.transformer = self.flux_pipe.transformer`
`38`		`- self.transformer.gradient_checkpointing = gradient_checkpointing`
	`38`	`+ if gradient_checkpointing:`
	`39`	`+ self.transformer.enable_gradient_checkpointing()`
`39`	`40`	`self.transformer.train()`
`40`	`41`
`41`	`42`	`# Freeze the Flux pipeline components`