diff --git a/train.py b/train.py
index 6d93268..0b35061 100644
--- a/train.py
+++ b/train.py
@@ -12,7 +12,7 @@
 from dataset_mask_val import Dataset as Dataset_val
 import os
 import torch
-from network import Res_Deeplab
+from one_shot_network import Res_Deeplab
 import torch.nn as nn
 import numpy as np
 
@@ -67,16 +67,17 @@
 options = parser.parse_args()
 
 
-data_dir = '/your/dataset/dir/VOCdevkit/VOC2012'
+data_dir = 'data'
 
 
 
 
 #set gpus
 gpu_list = [int(x) for x in options.gpu.split(',')]
+#print(gpu_list)
 os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
 os.environ['CUDA_VISIBLE_DEVICES'] = options.gpu
-
+#print(os.environ['CUDA_VISIBLE_DEVICES'])
 torch.backends.cudnn.benchmark = True
 
 
@@ -100,7 +101,7 @@
 model = Res_Deeplab(num_classes=num_class)
 #load resnet-50 preatrained parameter
 model = load_resnet50_param(model, stop_layer='layer4')
-model=nn.DataParallel(model,[0,1])
+#model=nn.DataParallel(model,[0,1])
 
 # disable the  gradients of not optomized layers
 turn_off(model)
@@ -136,8 +137,7 @@
 
 
 
-optimizer = optim.SGD([{'params': get_10x_lr_params(model), 'lr': 10 * learning_rate}],
-                          lr=learning_rate, momentum=momentum, weight_decay=weight_decay)
+optimizer = optim.SGD([{'params': get_10x_lr_params(model), 'lr': 10 * learning_rate}], lr=learning_rate, momentum=momentum, weight_decay=weight_decay)
 
 
 
diff --git a/utils.py b/utils.py
index 0d47ccf..8c97fa0 100644
--- a/utils.py
+++ b/utils.py
@@ -40,10 +40,10 @@ def optim_or_not(model, yes):
 
 
 def turn_off(model):
-    optim_or_not(model.module.conv1, False)
-    optim_or_not(model.module.layer1, False)
-    optim_or_not(model.module.layer2, False)
-    optim_or_not(model.module.layer3, False)
+    optim_or_not(model.conv1, False)
+    optim_or_not(model.layer1, False)
+    optim_or_not(model.layer2, False)
+    optim_or_not(model.layer3, False)
 
 
 
@@ -54,18 +54,18 @@ def get_10x_lr_params(model):
     """
 
     b = []
-    b.append(model.module.layer5.parameters())
-    b.append(model.module.layer55.parameters())
-    b.append(model.module.layer6_0.parameters())
-    b.append(model.module.layer6_1.parameters())
-    b.append(model.module.layer6_2.parameters())
-    b.append(model.module.layer6_3.parameters())
-    b.append(model.module.layer6_4.parameters())
-    b.append(model.module.layer7.parameters())
-    b.append(model.module.layer9.parameters())
-    b.append(model.module.residule1.parameters())
-    b.append(model.module.residule2.parameters())
-    b.append(model.module.residule3.parameters())
+    b.append(model.layer5.parameters())
+    b.append(model.layer55.parameters())
+    b.append(model.layer6_0.parameters())
+    b.append(model.layer6_1.parameters())
+    b.append(model.layer6_2.parameters())
+    b.append(model.layer6_3.parameters())
+    b.append(model.layer6_4.parameters())
+    b.append(model.layer7.parameters())
+    b.append(model.layer9.parameters())
+    b.append(model.residule1.parameters())
+    b.append(model.residule2.parameters())
+    b.append(model.residule3.parameters())
 
     for j in range(len(b)):
         for i in b[j]:
@@ -129,4 +129,4 @@ def get_iou_v1(query_mask,pred_label,mode='foreground'):#pytorch 1.0 version
             union_list.append(union)
             iou_list.append(0)
             num_predict_list.append(num_predict)
-    return inter_list,union_list,iou_list,num_predict_list
\ No newline at end of file
+    return inter_list,union_list,iou_list,num_predict_list