davidberenstein1957 commited on
Commit
724679b
·
verified ·
1 Parent(s): 53e269a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -12,7 +12,7 @@ device = "cuda" if torch.cuda.is_available() else "CPU"
12
 
13
  pipe = PrunaModel.from_hub(
14
  "PrunaAI/Segmind-Vega-smashed"
15
- )
16
 
17
  MAX_SEED = np.iinfo(np.int32).max
18
  MAX_IMAGE_SIZE = 2048
 
12
 
13
  pipe = PrunaModel.from_hub(
14
  "PrunaAI/Segmind-Vega-smashed"
15
+ ).move_to_device("cuda")
16
 
17
  MAX_SEED = np.iinfo(np.int32).max
18
  MAX_IMAGE_SIZE = 2048