cbensimon HF Staff commited on
Commit
060b6b6
·
1 Parent(s): 66ca860

Move weights to CUDA

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -53,6 +53,8 @@ def compile_model():
53
  return weights_
54
 
55
  weights = compile_model()
 
 
56
  del model
57
 
58
  @spaces.GPU
 
53
  return weights_
54
 
55
  weights = compile_model()
56
+ weights = {name: tensor.to('cuda') for name, tensor in weights.items()}
57
+
58
  del model
59
 
60
  @spaces.GPU