[rocm6.4_internal_testing] Temporary disable NHWC batchnorm tests in ROCm6.4 (#1849)

Remove definition of PYTORCH_MIOPEN_SUGGEST_NHWC environment variable to
disable NHWC batchnorm on MIOpen in ROCm6.4.
NHWC can be enabled back by defining PYTORCH_MIOPEN_SUGGEST_NHWC=1 if
needed
Skip related NHWC batchnorm tests
Waiting for NHWC batchnorm support from MIOpen (SWDEV-510757)

@pruthvistony should I remove NHWC batchnorm completely or is it enough
to disable it by default?

Author: dnikolaev-amd

test/test_nn.py

@@ -10,7 +10,6 @@
 import warnings
 import pickle
 import re
-import os
 from copy import deepcopy
 from itertools import product
 from functools import partial
@@ -5098,7 +5097,7 @@ def test_batchnorm_nhwc_eval(self, mixed, dtype):
         if TEST_WITH_ROCM and not mixed and dtype in (torch.half, torch.bfloat16):
             self.skipTest("pure mode not supported for bf16/fp16 on ROCm")
         if TEST_WITH_ROCM:
-            self.skipTest("MIOpen SolverNotFound for FP32/Fp16/BF16 NHWC batchnorm SWDEV-509640")
+            self.skipTest("NHWC batchnorm disabled on ROCm6.4 SWDEV-510757 SWDEV-509640")
 
         (N, C, H, W) = 2, 64, 50, 50
         model = torch.nn.BatchNorm2d(C, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
@@ -5159,6 +5158,8 @@ def _run_test(input, grad_output, mixed: bool):
             grad = grad.permute(0, 2, 1, 3)
             _run_test(input, grad, mixed)
 
+        if TEST_WITH_ROCM and layout == "NHWC":
+            self.skipTest("NHWC batchnorm disabled on ROCm6.4 SWDEV-510757")
         if mixed and dtype == torch.float:
             self.skipTest("mixed precision is useless for float32")
         if TEST_WITH_ROCM and not mixed and dtype in (torch.half, torch.bfloat16):
@@ -5167,17 +5168,7 @@ def _run_test(input, grad_output, mixed: bool):
             self.skipTest("MIOpen tolerance issue for NCHW BF16 mixed batchnorm SWDEV-507600")
 
         memory_format = torch.contiguous_format if layout == "NCHW" else torch.channels_last
-        # TODO: Remove PYTORCH_MIOPEN_SUGGEST_NHWC once ROCm officially supports NHWC in MIOpen
-        PYTORCH_MIOPEN_SUGGEST_NHWC = "PYTORCH_MIOPEN_SUGGEST_NHWC"
-        prev_val = os.getenv(PYTORCH_MIOPEN_SUGGEST_NHWC)
-        try:
-            os.environ[PYTORCH_MIOPEN_SUGGEST_NHWC] = "1"
-            _batchnorm2d_helper(dtype, memory_format=memory_format, mixed=mixed)
-        finally:
-            if prev_val is None:
-                del os.environ[PYTORCH_MIOPEN_SUGGEST_NHWC]
-            else:
-                os.environ[PYTORCH_MIOPEN_SUGGEST_NHWC] = prev_val
+        _batchnorm2d_helper(dtype, memory_format=memory_format, mixed=mixed)
 
     def test_batchnorm_load_state_dict(self):
         bn = torch.nn.BatchNorm2d(3)
@@ -13203,14 +13194,4 @@ def __init__(self) -> None:
 
 if __name__ == '__main__':
     TestCase._default_dtype_check_enabled = True
-    # TODO: Remove PYTORCH_MIOPEN_SUGGEST_NHWC once ROCm officially supports NHWC in MIOpen
-    PYTORCH_MIOPEN_SUGGEST_NHWC = "PYTORCH_MIOPEN_SUGGEST_NHWC"
-    prev_val = os.getenv(PYTORCH_MIOPEN_SUGGEST_NHWC)
-    try:
-        os.environ[PYTORCH_MIOPEN_SUGGEST_NHWC] = "1"
-        run_tests()
-    finally:
-        if prev_val is None:
-            del os.environ[PYTORCH_MIOPEN_SUGGEST_NHWC]
-        else:
-            os.environ[PYTORCH_MIOPEN_SUGGEST_NHWC] = prev_val
+    run_tests()