device check

tomrunia · Dec 14, 2018 · e54981e · e54981e
1 parent 0b6514d
commit e54981e
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/steerable/SCFpyr_PyTorch.py b/steerable/SCFpyr_PyTorch.py
@@ -74,6 +74,8 @@ def build(self, im_batch):
         Returns:
             pyramid: list containing torch.Tensor objects storing the pyramid
         '''
+
+        assert im_batch.device == self.device, 'Devices invalid (pyr = {}, batch = {})'.format(self.device, im_batch.device)
         assert im_batch.dtype == torch.float32, 'Image batch must be torch.float32'
         assert im_batch.dim() == 4, 'Image batch must be of shape [N,C,H,W]'
         assert im_batch.shape[1] == 1, 'Second dimension must be 1 encoding grayscale image'