Merge pull request #58 from fidelity/dataloader

Bug fix for moving data to gpu
fidelity · Jul 8, 2022 · 6f56534 · 6f56534
2 parents 2647afd + a263aa3
commit 6f56534
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/stoke/data.py b/stoke/data.py
@@ -105,7 +105,7 @@ def _place_data_on_gpu(
 
         """
         # Call to the public function in utils
-        place_data_on_gpu(data, self._fp16)
+        return place_data_on_gpu(data, self._fp16)
 
 
 class BucketedDistributedSampler(Sampler[T_co]):