From 451db0c4bc12a6cf790cd303f159ec33914672a7 Mon Sep 17 00:00:00 2001
From: Chengji Yao <chengjiyao@google.com>
Date: Tue, 4 Feb 2025 01:21:01 +0000
Subject: [PATCH] change back file importer

---
 torchax/test/test_core_aten_ops.py | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/torchax/test/test_core_aten_ops.py b/torchax/test/test_core_aten_ops.py
index 6dc99d48cc1..e4c0372c38f 100644
--- 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,