diff --git a/perfmetrics/scripts/continuous_test/ml_tests/pytorch/v2_hns/dino/build.sh b/perfmetrics/scripts/continuous_test/ml_tests/pytorch/v2_hns/dino/build.sh
index 1d9e15460a..39dee19905 100755
--- a/perfmetrics/scripts/continuous_test/ml_tests/pytorch/v2_hns/dino/build.sh
+++ b/perfmetrics/scripts/continuous_test/ml_tests/pytorch/v2_hns/dino/build.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright 2023 Google LLC
+# Copyright 2024 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/perfmetrics/scripts/continuous_test/ml_tests/run_and_manage_test.sh b/perfmetrics/scripts/continuous_test/ml_tests/run_and_manage_test.sh
index 19278222e7..9c1841a2e2 100755
--- a/perfmetrics/scripts/continuous_test/ml_tests/run_and_manage_test.sh
+++ b/perfmetrics/scripts/continuous_test/ml_tests/run_and_manage_test.sh
@@ -159,7 +159,7 @@ then
   delete_existing_vm_and_create_new
   
   echo "Clone the gcsfuse repo on test VM"
-  sudo gcloud compute ssh $VM_NAME --zone $ZONE_NAME --internal-ip --command "mkdir github; cd github; git clone https://github.com/GoogleCloudPlatform/gcsfuse.git; cd gcsfuse; git checkout ai_ml_tests_on_hns_bucket;"
+  sudo gcloud compute ssh $VM_NAME --zone $ZONE_NAME --internal-ip --command "mkdir github; cd github; git clone https://github.com/GoogleCloudPlatform/gcsfuse.git; cd gcsfuse; git checkout master;"
   echo "Trigger the build script on test VM"
   sudo gcloud compute ssh $VM_NAME --zone $ZONE_NAME --internal-ip --command "bash \$HOME/$TEST_SCRIPT_PATH 1> \$HOME/build.out 2> \$HOME/build.err &"
   echo "Wait for 15 minutes for test VM to setup for test and to change the status from START to RUNNING."
diff --git a/perfmetrics/scripts/continuous_test/ml_tests/tf/resnet/build.sh b/perfmetrics/scripts/continuous_test/ml_tests/tf/resnet/build.sh
index 8ff102e9ce..b98a908d7e 100755
--- a/perfmetrics/scripts/continuous_test/ml_tests/tf/resnet/build.sh
+++ b/perfmetrics/scripts/continuous_test/ml_tests/tf/resnet/build.sh
@@ -18,10 +18,9 @@ set -e
 
 VM_NAME="tf-resnet-7d"
 ZONE_NAME="us-west1-b"
-ARTIFACTS_BUCKET_PATH="gs://gcsfuse-ml-tests-logs/ci_artifacts/tf/resnet_hns"
-TEST_SCRIPT_PATH="github/gcsfuse/perfmetrics/scripts/ml_tests/tf/resnet_hns/setup_host_and_run_model.sh"
-BUCKET_TYPE="non-hns"
+ARTIFACTS_BUCKET_PATH="gs://gcsfuse-ml-tests-logs/ci_artifacts/tf/resnet"
+TEST_SCRIPT_PATH="github/gcsfuse/perfmetrics/scripts/ml_tests/tf/resnet/setup_host_and_run_model.sh"
 
 cd "${KOKORO_ARTIFACTS_DIR}/github/gcsfuse/perfmetrics/scripts/continuous_test/ml_tests/"
 
-source run_and_manage_test.sh $VM_NAME $ZONE_NAME $ARTIFACTS_BUCKET_PATH $TEST_SCRIPT_PATH $BUCKET_TYPE
+source run_and_manage_test.sh $VM_NAME $ZONE_NAME $ARTIFACTS_BUCKET_PATH $TEST_SCRIPT_PATH
diff --git a/perfmetrics/scripts/ml_tests/pytorch/run_container.sh b/perfmetrics/scripts/ml_tests/pytorch/run_container.sh
index f57280e4b8..8838484f8e 100644
--- a/perfmetrics/scripts/ml_tests/pytorch/run_container.sh
+++ b/perfmetrics/scripts/ml_tests/pytorch/run_container.sh
@@ -14,11 +14,11 @@
 # limitations under the License.
 set -e
 # pytorch version (e.g. v1_12, v2)
-PYTORCH_VERSION=$1
+DIR=$1
 
 cd "$HOME/github/gcsfuse"
 echo "Building docker image containing all pytorch libraries..."
-sudo docker build . -f perfmetrics/scripts/ml_tests/pytorch/${PYTORCH_VERSION}/dino/Dockerfile --tag pytorch-gcsfuse
+sudo docker build . -f perfmetrics/scripts/ml_tests/pytorch/${DIR}/dino/Dockerfile --tag pytorch-gcsfuse
 
 mkdir -p container_artifacts