Support bool dtype (#22)

tests/test_torchsnooper.py

@@ -314,3 +314,25 @@ def my_function():
  ReturnEntry(),
  )
  )
+
+
+def test_bool_tensor():
+ string_io = io.StringIO()
+
+ @torchsnooper.snoop(string_io)
+ def my_function():
+ x = torch.zeros(5, 5, dtype=torch.bool) # noqa: F841
+
+ my_function()
+
+ output = string_io.getvalue()
+ print(output)
+ assert_output(
+ output,
+ (
+ CallEntry(),
+ LineEntry(),
+ VariableEntry('x', "tensor<(), bool, cpu>"),
+ ReturnEntry(),
+ )
+ )

torchsnooper/__init__.py

@@ -6,6 +6,12 @@
 from pkg_resources import get_distribution, DistributionNotFound
 
 
+FLOATING_POINTS = set()
+for i in ['float', 'double', 'half', 'complex128', 'complex32', 'complex64']:
+ if hasattr(torch, i): # older version of PyTorch do not have complex dtypes
+ FLOATING_POINTS.add(getattr(torch, i))
+
+
 try:
  __version__ = get_distribution(__name__).version
 except DistributionNotFound:
@@ -47,15 +53,15 @@ def __call__(self, tensor):
  if tensor.requires_grad:
  new += 'grad'
  elif p == 'has_nan':
- result = bool(torch.isnan(tensor).any())
+ result = tensor.dtype in FLOATING_POINTS and bool(torch.isnan(tensor).any())
  if self.properties_name:
  new += 'has_nan='
  new += str(result)
  else:
  if result:
  new += 'has_nan'
  elif p == 'has_inf':
- result = bool(torch.isinf(tensor).any())
+ result = tensor.dtype in FLOATING_POINTS and bool(torch.isinf(tensor).any())
  if self.properties_name:
  new += 'has_inf='
  new += str(result)