def collate_fn()

in jcm/datasets.py [0:0]

• 41 lines of code
• 4 McCabe index (conditional complexity)

        def collate_fn(batch):
            image_tensor = torch.stack([x[0] for x in batch], dim=0)
            label_tensor = torch.tensor([x[1] for x in batch])
            if image_tensor.shape[0] == int(np.prod(batch_dims)):
                image_tensor = (
                    image_tensor.reshape(batch_dims + [3, 32, 32])
                    .transpose(-3, -2)
                    .transpose(-2, -1)
                )
                label_tensor = label_tensor.reshape(batch_dims)
                return {
                    "image": image_tensor,
                    "label": label_tensor,
                    "mask": torch.ones_like(label_tensor),
                }
            # If the batch size is not a multiple of the batch dimension, pad the batch with zeros.
            else:
                pad_size = int(np.prod(batch_dims)) - image_tensor.shape[0]
                padded_image = torch.concat(
                    [
                        image_tensor,
                        torch.zeros(pad_size, 3, 32, 32, dtype=image_tensor.dtype),
                    ],
                    axis=0,
                )
                padded_label = torch.concat(
                    [
                        label_tensor,
                        torch.zeros(pad_size, dtype=label_tensor.dtype),
                    ],
                    axis=0,
                )
                mask = torch.ones(int(np.prod(batch_dims)))
                mask[image_tensor.shape[0] :] = 0.0
                padded_image = (
                    padded_image.reshape(batch_dims + [3, 32, 32])
                    .transpose(-3, -2)
                    .transpose(-2, -1)
                )
                padded_label = padded_label.reshape(batch_dims)
                mask = mask.reshape(batch_dims)
                return {"image": padded_image, "label": padded_label, "mask": mask}