Make it PT Only compatible

Author: Ubuntu

docker/Dockerfile

@@ -46,12 +46,12 @@ RUN if [ "${cloud_build}" = true ]; then github_branch="${release_version}" && \
   cd / && \
   git clone -b "${github_branch}" --recursive https://github.com/pytorch-tpu/examples tpu-examples; fi
 
-# RUN cd /pytorch && bash xla/scripts/build_torch_wheels.sh ${python_version} ${release_version}
+RUN cd /pytorch && bash xla/scripts/build_torch_wheels.sh ${python_version} ${release_version}
 
 # Use conda environment on startup or when running scripts.
-# RUN echo "conda activate pytorch" >> ~/.bashrc
-# RUN echo "export TF_CPP_LOG_THREAD_ID=1" >> ~/.bashrc
-# ENV PATH /root/anaconda3/envs/pytorch/bin/:/root/bin:$PATH
+RUN echo "conda activate pytorch" >> ~/.bashrc
+RUN echo "export TF_CPP_LOG_THREAD_ID=1" >> ~/.bashrc
+ENV PATH /root/anaconda3/envs/pytorch/bin/:/root/bin:$PATH
 
 # Define entrypoint and cmd
 COPY docker/docker-entrypoint.sh /usr/local/bin

test/__init__.py

@@ -1,4 +1,4 @@
-import args_parse
+import test.args_parse as args_parse
 
 SUPPORTED_MODELS = [
     "alexnet",
@@ -56,7 +56,7 @@
 
 import os
 
-import schedulers
+import test.schedulers as schedulers
 import numpy as np
 import torch
 import torch.nn as nn
@@ -67,7 +67,7 @@
 import torch_xla.utils.utils as xu
 import torch_xla.test.test_utils as test_utils
 import torch_xla.core.xla_model as xm
-from classification_benchmark_constants import DEFAULT_KWARGS, MODEL_SPECIFIC_DEFAULTS
+from test.classification_benchmark_constants import DEFAULT_KWARGS, MODEL_SPECIFIC_DEFAULTS
 from torch.cuda.amp import GradScaler, autocast
 
 default_value_dict = MODEL_SPECIFIC_DEFAULTS.get(FLAGS.model, DEFAULT_KWARGS)