def collate_batch(self, batch):
b = collate_batch(batch)
for i in range(len(self.args)):
x = b[i]
if jt.is_var(self.args[i]) and self.args[i].ndim == 1:
x.assign(x.squeeze(-1))
return b
def collate_batch(self, batch):
'''
Puts each data field into a tensor with outer dimension batch size.
Args::
[in] batch(list): A list of variables, such as jt.var, Image.Image, np.ndarray, int, float, str and so on.
'''
return collate_batch(batch)
def collate_batch(self, batch):
return collate_batch(batch)
def test_collate_batch(self):
from jittor.dataset.utils import collate_batch
batch = collate_batch([(1, 1), (1, 2), (1, 3)])
assert isinstance(batch[0], np.ndarray)
assert isinstance(batch[1], np.ndarray)
