change back file importer

pytorch · Feb 4, 2025 · 451db0c · 451db0c
1 parent 10c124d
commit 451db0c
Showing 1 changed file with 2 additions and 10 deletions.
diff --git a/torchax/test/test_core_aten_ops.py b/torchax/test/test_core_aten_ops.py
@@ -1,19 +1,11 @@
 import math
-import os
-import sys
 import unittest
 
-
 import torch
-from torch.utils import _pytree as pytree
 from torchax import tensor
 
-# Setup import folders.
-test_folder = os.path.dirname(os.path.abspath(sys.argv[0]))
-sys.path.append(test_folder)
-
-import test_base
-
+from . import test_base
+from torch.utils import _pytree as pytree
 
 
 def diff_output(testcase,