def _after_forward(self, **kwargs):
        """
        Reshape the model output to have 2 views: one for original examples,
        one for augmented examples.
        """
        assert self.is_training
        super()._after_forward(**kwargs)
        assert self.mb_output.size(0) % 2 == 0
        original_batch_size = int(self.mb_output.size(0) / 2)
        original_examples = self.mb_output[:original_batch_size]
        augmented_examples = self.mb_output[original_batch_size:]
        # (original_batch_size, 2, output_size)
        self.mb_output = torch.stack([original_examples, augmented_examples], dim=1)

        Reshape the model output to have 2 views: one for original examples,
        one for augmented examples.
        
avalanche/training/supervised/supervised_contrastive_replay.py