take_along_axis validation error (race condition?)

[naoyam]

Reported by @csarofeen.

[==========] Running 1 test from 1 test suite.
[----------] Global test environment set-up.
[----------] 1 test from ScatterGatherTest
[ RUN      ] ScatterGatherTest.TakeAlongAxisIntermediateTensorNormalizationAndReduction2
unknown file: Failure
C++ exception with description " INTERNAL ASSERT FAILED at "/opt/pytorch/Fuser/tests/cpp/validator.cpp":115, please report a bug with repro script to NVFuser at https://github.com/NVIDIA/Fuser/issues.

Validation error in output 0 on line 994 in file /opt/pytorch/Fuser/tests/cpp/test_scatter_gather.cpp.
  Detected max abs error of: 189.495
    absolute tolerance was set to 0.00343043
    and relative tolerance set to 3.43043e-05
Exception raised from testValidate at /opt/pytorch/Fuser/tests/cpp/validator.cpp:115 (most recent call first):
frame #0: <unknown function> + 0x1c5300 (0xc89c41405300 in ./bin/test_nvfuser)
frame #1: <unknown function> + 0x46cbf0 (0xc89c416acbf0 in ./bin/test_nvfuser)
frame #2: <unknown function> + 0x1116858 (0xc89c42356858 in ./bin/test_nvfuser)
frame #3: <unknown function> + 0x1055ad0 (0xc89c42295ad0 in ./bin/test_nvfuser)
frame #4: <unknown function> + 0x1153b10 (0xc89c42393b10 in ./bin/test_nvfuser)
frame #5: <unknown function> + 0x113be24 (0xc89c4237be24 in ./bin/test_nvfuser)
frame #6: <unknown function> + 0x113c318 (0xc89c4237c318 in ./bin/test_nvfuser)
frame #7: <unknown function> + 0x113c914 (0xc89c4237c914 in ./bin/test_nvfuser)
frame #8: <unknown function> + 0x1149bc0 (0xc89c42389bc0 in ./bin/test_nvfuser)
frame #9: <unknown function> + 0x113caf0 (0xc89c4237caf0 in ./bin/test_nvfuser)
frame #10: <unknown function> + 0x223700 (0xc89c41463700 in ./bin/test_nvfuser)
frame #11: <unknown function> + 0x284c4 (0xe2373de584c4 in /usr/lib/aarch64-linux-gnu/libc.so.6)
frame #12: __libc_start_main + 0x98 (0xe2373de58598 in /usr/lib/aarch64-linux-gnu/libc.so.6)
frame #13: <unknown function> + 0x223d70 (0xc89c41463d70 in ./bin/test_nvfuser)
" thrown in the test body.

To reproduce: NVFUSER_TEST_RANDOM_SEED=1741024131 NVFUSER_TEST_ATEN_RANDOM_SEED=0 test_nvfuser --gtest_filter='ScatterGatherTest.TakeAlongAxisIntermediateTensorNormalizationAndReduction2'
[  FAILED  ] ScatterGatherTest.TakeAlongAxisIntermediateTensorNormalizationAndReduction2 (487 ms)
[----------] 1 test from ScatterGatherTest (487 ms total)

[----------] Global test environment tear-down
[==========] 1 test from 1 test suite ran. (487 ms total)
[  PASSED  ] 0 tests.
[  FAILED  ] 1 test, listed below:
[  FAILED  ] ScatterGatherTest.TakeAlongAxisIntermediateTensorNormalizationAndReduction2