983a888f498131911a1296400448b004a53c1717,pytorch_lightning/accelerators/tpu.py,TPUAccelerator,all_gather,#TPUAccelerator#Any#Any#Any#,44

Before Change

        
        // todo: Add support for backward with all_gather
        if torch.distributed.is_initialized():
            return xm.all_gather(tensor, group=group, sync_grads=sync_grads)
        return tensor

After Change

        
        // todo: Add support for backward with all_gather
        if isinstance(self.training_type_plugin, TPUSpawnPlugin) and self.training_type_plugin.is_distributed:
            return xm.all_gather(tensor).view(-1, *tensor.shape)
        return tensor

Instances