make MXDataIter work without indices (#7456)

indices are optional, custom cpp iterators providing data batches without indices should work while using MXDataIter.
apache · Aug 24, 2017 · 4b94360 · 4b94360
1 parent b710636
commit 4b94360
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/python/mxnet/io.py b/python/mxnet/io.py
@@ -815,10 +815,13 @@ def getindex(self):
         check_call(_LIB.MXDataIterGetIndex(self.handle,
                                            ctypes.byref(index_data),
                                            ctypes.byref(index_size)))
-        address = ctypes.addressof(index_data.contents)
-        dbuffer = (ctypes.c_uint64* index_size.value).from_address(address)
-        np_index = np.frombuffer(dbuffer, dtype=np.uint64)
-        return np_index.copy()
+        if index_size.value:
+            address = ctypes.addressof(index_data.contents)
+            dbuffer = (ctypes.c_uint64* index_size.value).from_address(address)
+            np_index = np.frombuffer(dbuffer, dtype=np.uint64)
+            return np_index.copy()
+        else:
+            return None
 
     def getpad(self):
         pad = ctypes.c_int(0)