def METHOD_NAME(self):
    losses = torch.rand(2, 2).t()
    gathered_loss = self.all_gather(
        {
            "losses_tensor_int": losses.int(),
            "losses_tensor_float": losses,
            "losses_tensor_list": [losses, losses],
            "losses_np_ndarray": np.array([1, 2, 3]),
            "losses_bool": [True, False],
            "losses_float": [0.0, 1.0, 2.0],
            "losses_int": [0, 1, 2],
        }
    )
    assert gathered_loss["losses_tensor_int"][0].dtype == torch.int32
    assert gathered_loss["losses_tensor_float"][0].dtype == torch.float
    assert gathered_loss["losses_np_ndarray"][0].dtype == torch.int64
    # torch.bool can't be all_gathered
    assert gathered_loss["losses_bool"][0].dtype == torch.uint8
    assert gathered_loss["losses_float"][0].dtype == torch.float
    assert gathered_loss["losses_int"][0].dtype == torch.int
    losses_numel = losses.numel()
    assert gathered_loss["losses_tensor_int"].numel() == 2 * losses_numel
    assert gathered_loss["losses_tensor_float"].numel() == 2 * losses_numel
    assert torch.stack(gathered_loss["losses_tensor_list"]).shape == (2, 2, 2, 2)
    assert gathered_loss["losses_np_ndarray"].numel() == 2 * 3
    assert torch.stack(gathered_loss["losses_bool"]).shape == (2, 2)
    assert torch.stack(gathered_loss["losses_float"]).shape == (3, 2)
    assert torch.stack(gathered_loss["losses_int"]).shape == (3, 2)
    self.on_train_epoch_end_called = True