Accept longs in default_collate for dataloader in python 2 (#4001)

alykhantejani · soumith · commit e98af60a7dca · 2018-02-06T20:51:04.000-08:00
diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py
@@ -6,7 +6,7 @@
 import sys
 import traceback
 import threading
-from torch._six import string_classes
+from torch._six import string_classes, int_classes
 
 
 if sys.version_info[0] == 2:
@@ -106,7 +106,7 @@ def default_collate(batch):
         if elem.shape == ():  # scalars
             py_type = float if elem.dtype.name.startswith('float') else int
             return numpy_type_map[elem.dtype.name](list(map(py_type, batch)))
-    elif isinstance(batch[0], int):
+    elif isinstance(batch[0], int_classes):
         return torch.LongTensor(batch)
     elif isinstance(batch[0], float):
         return torch.DoubleTensor(batch)
