Spaces:

rawc0der
/

t2i-model-mayhem

Sleeping

rawc0der commited on Jan 19

Commit

34a68a5

1 Parent(s): 96db79f

fix dispatch

Files changed (1) hide show

app.py CHANGED Viewed

@@ -161,7 +161,7 @@ class ModelContext:
         if hasattr(pipeline, 'reset_device_map'):
             pipeline.reset_device_map()
         # Enable automatic device mapping strategy
-        dispatch_model(pipeline)
         return pipeline
     def __exit__(self, exc_type, exc_val, exc_tb):

         if hasattr(pipeline, 'reset_device_map'):
             pipeline.reset_device_map()
         # Enable automatic device mapping strategy
+        dispatch_model(pipeline, device_map="auto")
         return pipeline
     def __exit__(self, exc_type, exc_val, exc_tb):