Import missing packages

Author: ghpvnist

test/pjrt/test_dtypes.py

@@ -1,5 +1,6 @@
 from absl.testing import absltest, parameterized
 import torch
+import torch_xla
 import torch_xla.core.xla_model as xm
 import torch_xla.runtime as xr
 

test/stablehlo/test_pt2e_qdq.py

@@ -4,6 +4,7 @@
 from typing import Callable, Dict, List
 
 import torch
+import torch_xla
 import torch_xla.core.xla_model as xm
 import torchvision
 from torch.export import export_for_training