FrankZxShen
commited on
Commit
•
9a8c488
1
Parent(s):
78121a7
Update inference/infer_tool.py
Browse files- inference/infer_tool.py +2 -2
inference/infer_tool.py
CHANGED
@@ -127,8 +127,8 @@ class Svc(object):
|
|
127 |
self.only_diffusion = only_diffusion
|
128 |
self.shallow_diffusion = shallow_diffusion
|
129 |
if device is None:
|
130 |
-
|
131 |
-
self.dev = torch.device("cpu")
|
132 |
else:
|
133 |
self.dev = torch.device(device)
|
134 |
self.net_g_ms = None
|
|
|
127 |
self.only_diffusion = only_diffusion
|
128 |
self.shallow_diffusion = shallow_diffusion
|
129 |
if device is None:
|
130 |
+
self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
131 |
+
# self.dev = torch.device("cpu")
|
132 |
else:
|
133 |
self.dev = torch.device(device)
|
134 |
self.net_g_ms = None
|