fix dataset to_jittor bug

2020-12-16 15:56:34 +08:00 · 2020-12-16 15:56:34 +08:00 · 23de860d68
parent 73eb05b36e
commit 23de860d68
1 changed files with 1 additions and 0 deletions
--- a/python/jittor/dataset/dataset.py
+++ b/python/jittor/dataset/dataset.py
@ -117,6 +117,7 @@ class Dataset(object):
        '''
        Change batch data to jittor array, such as np.ndarray, int, and float.
        '''
+        if isinstance(batch, jt.Var): return batch
        to_jt = lambda x: jt.array(x).stop_grad() \
            if self.stop_grad else jt.array(x)
        if isinstance(batch, np.ndarray):