methods/utils/st/nnst: update utils for nnst method

deel-ai · Jan 31, 2024 · 0bd791b · 0bd791b
1 parent 9438844
commit 0bd791b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/augmentare/methods/utils/style_transfer/utils_nnst.py b/augmentare/methods/utils/style_transfer/utils_nnst.py
@@ -283,7 +283,8 @@ def replace_features(src, ref):
     """
     # Move style features to gpu (necessary to mostly store on cpu for gpus w/
     # < 12GB of memory)
-    ref_flat = to_device(flatten_grid(ref))
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    ref_flat = flatten_grid(ref).to(device)
     rplc = []
     for j in range(src.size(0)):
         # How many rows of the distance matrix to compute at once, can be