Spaces:
Runtime error
Runtime error
Jaemin Cho commited on
Commit Β·
6fc6364
1
Parent(s): 46fd0bc
Use gpu
Browse files
app.py
CHANGED
|
@@ -52,6 +52,10 @@ class Instance:
|
|
| 52 |
|
| 53 |
print("Loaded model")
|
| 54 |
|
|
|
|
|
|
|
|
|
|
|
|
|
| 55 |
# # This command loads the individual model components on GPU on-demand. So, we don't
|
| 56 |
# # need to explicitly call pipe.to("cuda").
|
| 57 |
# pipe.enable_model_cpu_offload()
|
|
|
|
| 52 |
|
| 53 |
print("Loaded model")
|
| 54 |
|
| 55 |
+
if torch.cuda.is_available():
|
| 56 |
+
pipe.to("cuda")
|
| 57 |
+
print("Loaded model to GPU")
|
| 58 |
+
|
| 59 |
# # This command loads the individual model components on GPU on-demand. So, we don't
|
| 60 |
# # need to explicitly call pipe.to("cuda").
|
| 61 |
# pipe.enable_model_cpu_offload()
|