def __call__(self, batch: List[torch.Tensor]):
"""
:param batch: batch format [no_samples, height, width, channels]
:return:
"""
# batch contains a list of tuples of structure (sequence, target)
inputs = [item[0].to(self.device) for item in batch]
inputs = torch.stack(inputs)
# transform into vector
# inputs = inputs.view(inputs.shape[0], -1)
# transform into CHW matrix
# inputs = inputs.permute(0, 3, 1, 2)
inputs = inputs.reshape(-1, 1, 28, 28)
targets_tensor = torch.tensor([item[1]
for item in batch]).to(inputs[0].device)
target_partitions = {self.target_publication_key: targets_tensor}
return DatasetBatch(targets=target_partitions,
tags=None,
samples=inputs)
def dataset_batch(self) -> DatasetBatch:
tensor = torch.IntTensor([0, 0, 0, 1, 1, 1])
return DatasetBatch(
targets={TestDatasetBatch.target_key: tensor.clone()},
samples=tensor.clone(),
tags=tensor.clone())