cleaning up demos

2026-09-11 12:31:23 +08:00 · 2019-10-05 16:30:27 -04:00
parent dd85e1c600
commit 1c45edd192
1 changed files with 1 additions and 1 deletions
@@ -35,7 +35,7 @@ class ModelSummary(object):
        out_sizes = []
        input_ = self.model.example_input_array

-        if self.model.use_ddp or self.model.use_dp or self.model.single_gpu:
+        if self.model.on_gpu:
            input_ = input_.cuda(0)

        if self.model.trainer.use_amp: