manbeast3b
/

simpleflux6

Model card Files Files and versions

manbeast3b commited on Nov 28, 2024

Commit

afe4e6b

·

verified ·

1 Parent(s): 3238c79

Update src/pipeline.py

Files changed (1) hide show

src/pipeline.py +1 -1

src/pipeline.py CHANGED Viewed

@@ -17,7 +17,7 @@ def clear():
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()
 def load_pipeline() -> Pipeline:
     clear()
     dtype, device = torch.bfloat16, "cuda"

     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()
+@torch.inference_mode()
 def load_pipeline() -> Pipeline:
     clear()
     dtype, device = torch.bfloat16, "cuda"