adding a better assert error message

facebookresearch · Mar 30, 2021 · 2f4ce67 · 2f4ce67
1 parent 1d4ddbe
commit 2f4ce67
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/fairscale/nn/misc/param_bucket.py b/fairscale/nn/misc/param_bucket.py
@@ -73,8 +73,12 @@ def add_param(self, param: torch.Tensor) -> None:
     @torch.no_grad()
     def _add_param_as_view(self, param: torch.Tensor, keep_existing_value: bool = True) -> None:
         assert self.buffer is not None
-        assert param.dtype == self.buffer.dtype
-        assert param.device == self.buffer.device
+        assert (
+            param.dtype == self.buffer.dtype
+        ), f"Different types for the bucket and the param, cannot proceed: {param.dtype} - {self.buffer.dtype}"
+        assert (
+            param.device == self.buffer.device
+        ), f"Different devices for the bucket and the param, cannot proceed: {param.device} - {self.buffer.device}"
 
         fill_next = self._fill + param.numel()
         assert fill_next <= self.buffer.numel()

diff --git a/pyproject.toml b/pyproject.toml
@@ -28,4 +28,4 @@ use_parentheses = true
 skip_glob = ["build/*", "stubs/*"]
 # Don't split "import" and "from".
 force_sort_within_sections = true
-known_third_party = ["benchmark_dataset", "datasets", "golden_configs", "helpers", "models", "numpy", "parameterized", "pytest", "recommonmark", "setuptools", "torch", "torch_pg", "torchtext", "torchvision"]
+known_third_party = ["benchmark_dataset", "dataclasses", "datasets", "golden_configs", "helpers", "models", "numpy", "parameterized", "pytest", "recommonmark", "setuptools", "torch", "torch_pg", "torchtext", "torchvision"]