Adding a 64GB limit for cross_entropy_large test

- https://ontrack-internal.amd.com/browse/SWDEV-373709
  ROCm needs more memory and no GPU support with >64GB,
  so keeping this change on ROCm fork only.

Author: pruthvistony

test/test_nn.py

@@ -11554,7 +11554,8 @@ def check_equal(loss, inp_targ_1, inp_targ_2):
     # Ref: https://github.com/pytorch/pytorch/issue/85005
     @onlyCUDA
     @largeTensorTest("45GB", "cpu")
-    @largeTensorTest("45GB", "cuda")
+    # https://ontrack-internal.amd.com/browse/SWDEV-373709, ROCm needs more memory and no GPU support with >64GB
+    @largeTensorTest("64GB" if TEST_WITH_ROCM else "45GB", "cuda")
     @parametrize_test("reduction", ("none", "mean", "sum"))
     def test_cross_entropy_large_tensor(self, device, reduction):
         logits = torch.randn(int(2 ** 16), int(2 ** 16) + 1, dtype=torch.float32, device='cuda', requires_grad=True)