Merge pull request #16 from LukasHedegaard/develop

Fix missing device transfer in benchmark
LukasHedegaard · Mar 18, 2022 · b3aa999 · b3aa999
2 parents 20836c6 + cfe1fa8
commit b3aa999
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 
+## [0.7.1] - 2022-03-18
+
+### Fixed
+- Device transfer in benchmark.
+
+
 ## [0.7.0] - 2022-03-18
 
 ### Added

diff --git a/ride/info.py b/ride/info.py
@@ -1,6 +1,6 @@
 import time
 
-__version__ = "0.7.0"
+__version__ = "0.7.1"
 __author__ = "Lukas Hedegaard"
 __author_email__ = "lukasxhedegaard@gmail.com"
 __license__ = "Apache-2.0"

diff --git a/ride/runner.py b/ride/runner.py
@@ -202,6 +202,9 @@ def profile_model(self, args: AttributeDict, num_runs: int = 100) -> Dict[str, A
         else:
             model = self.Module(hparams=args)
 
+        if torch.cuda.is_available():
+            model = model.cuda()
+
         sample = torch.randn(
             getattr(model.hparams, "batch_size", 1), *model.input_shape
         )