diff --git a/gpu_mnist_example/train_gpu.py b/gpu_mnist_example/train_gpu.py
index f163b42..f83a473 100644
--- a/gpu_mnist_example/train_gpu.py
+++ b/gpu_mnist_example/train_gpu.py
@@ -12,7 +12,7 @@ If there are Chinese comments in the code，please add at the beginning：
 
 '''
 
-
+import os
 from model import Model
 import numpy as np
 import torch
@@ -22,10 +22,10 @@ from torch.optim import SGD
 from torch.utils.data import DataLoader
 from torchvision.transforms import ToTensor
 import argparse
-import os
 #导入openi包
 from openi.context import prepare, upload_openi
 
+os.system("pip install -U openi_test")
 # Training settings
 parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
 parser.add_argument('--epoch_size', type=int, default=10, help='how much epoch to train')