diff --git a/gpu_new/train.py b/gpu_new/train.py
index c6e3dfb..50da2a5 100755
--- a/gpu_new/train.py
+++ b/gpu_new/train.py
@@ -39,6 +39,11 @@ parser.add_argument('--testdata', default="/dataset/test" ,help='path to test da
 parser.add_argument('--epoch_size', type=int, default=1, help='how much epoch to train')
 parser.add_argument('--batch_size', type=int, default=256, help='how much batch_size in epoch')
 
+# 参数声明
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+model = Model().to(device)
+optimizer = SGD(model.parameters(), lr=1e-1)
+
 if __name__ == '__main__':
     args, unknown = parser.parse_known_args()
     #log output