def METHOD_NAME():
    """Test that casting to categorical works as expected."""
    test_tensor = torch.stack(
        [
            torch.cat([torch.eye(5, dtype=int), torch.zeros((5, 5), dtype=int)]),
            torch.cat([torch.zeros((5, 5), dtype=int), torch.eye(5, dtype=int)]),
        ]
    ).to(torch.float)
    expected = tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
    assert expected.shape == (2, 5)
    assert test_tensor.shape == (2, 10, 5)
    result = to_categorical(test_tensor)
    assert result.shape == expected.shape
    assert torch.allclose(result, expected.to(result.dtype))