Introducing IsInf

docs/source/torch.rst

@@ -222,6 +222,7 @@ Comparison Ops
 .. autofunction:: equal
 .. autofunction:: ge
 .. autofunction:: gt
+.. autofunction:: isinf
 .. autofunction:: isnan
 .. autofunction:: kthvalue
 .. autofunction:: le

test/test_torch.py

@@ -4533,6 +4533,10 @@ def test_isnan(self):
         x = torch.Tensor([1, float('nan'), 2])
         self.assertEqual(torch.isnan(x), torch.ByteTensor([0, 1, 0]))
 
+    def test_isinf(self):
+        x = torch.Tensor([1, float('inf'), 2, float('-inf'), float('nan')])
+        self.assertEqual(torch.isinf(x), torch.ByteTensor([0, 1, 0, 1, 0]))
+
     def test_RNGState(self):
         state = torch.get_rng_state()
         stateCloned = state.clone()

torch/functional.py

@@ -8,6 +8,7 @@
     'argmin',
     'btrifact',
     'btriunpack',
+    'isinf',
     'isnan',
     'split',
     'unbind',
@@ -148,6 +149,25 @@ def btriunpack(LU_data, LU_pivots, unpack_data=True, unpack_pivots=True):
     return P, L, U
 
 
+def isinf(tensor):
+    r"""Returns a new tensor with boolean elements representing if each element is `+/-INF` or not.
+
+    Arguments:
+        tensor (Tensor): A tensor to check
+
+    Returns:
+        Tensor: A ``torch.ByteTensor`` containing a 1 at each location of `+/-INF` elements and 0 otherwise
+
+    Example::
+
+        >>> torch.isinf(torch.Tensor([1, float('inf'), 2, float('-inf'), float('nan')]))
+        tensor([ 0,  1,  0,  1,  0], dtype=torch.uint8)
+    """
+    if not isinstance(tensor, torch.Tensor):
+        raise ValueError("The argument is not a tensor", str(tensor))
+    return tensor.abs() == float('inf')
+
+
 def isnan(tensor):
     r"""Returns a new tensor with boolean elements representing if each element is `NaN` or not.
 
@@ -163,7 +183,7 @@ def isnan(tensor):
         tensor([ 0,  1,  0], dtype=torch.uint8)
     """
     if not isinstance(tensor, torch.Tensor):
-        raise ValueError("The argument is not a tensor")
+        raise ValueError("The argument is not a tensor", str(tensor))
     return tensor != tensor