NPU implementation for FLUX

蒋硕 · 蒋硕 · commit 53f3c64c3455 · 2024-10-24T20:18:22.000+08:00
diff --git a/examples/dreambooth/train_dreambooth_flux.py b/examples/dreambooth/train_dreambooth_flux.py
@@ -193,6 +193,8 @@ def log_validation(
     del pipeline
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+    elif is_torch_npu_available():
+        torch_npu.npu.empty_cache()
 
     return images
 
