diff --git a/.github/workflows/PR_CI.yml b/.github/workflows/PR_CI.yml
index 9e2658973..7dccfbecd 100644
--- a/.github/workflows/PR_CI.yml
+++ b/.github/workflows/PR_CI.yml
@@ -28,7 +28,7 @@ jobs:
     runs-on: [self-hosted, gpu]
     container:
       image: nvcr.io/nvidia/pytorch:21.07-py3
-      options: --gpus all --rm --ipc=host -v /data/cifar-10:/data/cifar-10
+      options: --gpus all --rm --ipc=host -v /data/scratch/cifar-10:/data/scratch/cifar-10
     timeout-minutes: 20
     steps:
       - name: Setup Environment
@@ -48,7 +48,7 @@ jobs:
         run: |
           pytest tests
         env:
-          DATA: /data/cifar-10
+          DATA: /data/scratch/cifar-10
 
   format_check:
     name: Format Check
diff --git a/colossalai/nn/metric/accuracy_2p5d.py b/colossalai/nn/metric/accuracy_2p5d.py
index 0eeedd46f..891594264 100644
--- a/colossalai/nn/metric/accuracy_2p5d.py
+++ b/colossalai/nn/metric/accuracy_2p5d.py
@@ -1,5 +1,5 @@
 import torch
-from colossalai.nn.layer.parallel_2p5d import reduce_by_batch_2p5d
+from colossalai.nn.layer.parallel_2p5d import reduce_by_batch_2p5d, split_tensor_2p5d
 from torch import nn
 
 from ._utils import calc_acc