fixed bug on cpu-only machine

simplymathematics · Dec 1, 2023 · 8f47796 · 8f47796
1 parent 13c7852
commit 8f47796
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/deckard/base/model/art_pipeline.py b/deckard/base/model/art_pipeline.py
@@ -77,6 +77,8 @@ def __call__(self):
                     data = [torch.from_numpy(d).to(device) for d in data]
                 data = [d.to(device) for d in data]
                 model.to(device)
+            else:
+                device = torch.device("cpu")
             model = TorchInitializer(
                 data=data,
                 model=model,