I am using dataparallel in my network to train my model, but for testing if i use the same model it gives me error (KeyError: 'unexpected key "module.Base.conv1_1.conv.0.weight" in state_dict').
to fix this i had to add model = torch.nn.DataParallel(model.cuda()) again in my code, which i dont understand why is that the reason.
Also if i uncomment the part that i need for testing and want to train in that situation,  i get the error: AttributeError: 'DataParallel' object has no attribute 'init_parameters'
Can anyone explain it to me, or if im doing something wrong i appreciate if someone tell me what im doing wrong.
I provide  the code for training/testing part here, sorry it is kinda long, i comment the part that i have to uncomment for testing in this code and it has a comment (“Uncomment for test”) on the top of it.
Thanks for your help in advance

import os
import time
import argparse
import random
import numpy as np
from numpy.random import RandomState
import cv2
import torch.backends.cudnn as cudnn

import torch
import torch.optim as optim
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
from torch.utils.data import DataLoader

from pascal_voc import VOCDetection, VOC, Viz
from transforms import *
from evaluation import eval_voc_detection

import sys
sys.path.append('./models')
from SSD import SSD300

from loss import MultiBoxLoss
from multibox import MultiBox


from tensorboardX import SummaryWriter
summary = SummaryWriter()

# Setup
parser = argparse.ArgumentParser(description='PyTorch SSD variants implementation')
parser.add_argument('--checkpoint', help='resume from checkpoint', default='')
parser.add_argument('--voc_root', help='PASCAL VOC dataset root path', default='')
parser.add_argument('--batch_size', type=int, help='input data batch size', default=32)
parser.add_argument('--lr', '--learning_rate', type=float, help='initial learning rate', default=1e-3)
parser.add_argument('--start_iter', type=int, help='start iteration', default=0)
parser.add_argument('--backbone', help='pretrained backbone net weights file', default='vgg16_reducedfc.pth')
parser.add_argument('--cuda', action='store_true', help='enable cuda')
parser.add_argument('--test', action='store_true', help='test mode')
parser.add_argument('--demo', action='store_true', help='show detection result')
parser.add_argument('--seed', type=int, help='random seed', default=233)
parser.add_argument('--threads', type=int, help='number of data loader workers.', default=4)
parser.add_argument('--max_iter', type=int, help='max_iter number.', default=100000)
parser.add_argument('--width_mult', type=int, help='width_mult number.', default=1)

opt = parser.parse_args()
print('argparser:', opt)




# random seed
random.seed(opt.seed)
np.random.seed(opt.seed)
torch.manual_seed(opt.seed)
if opt.cuda:
    assert torch.cuda.is_available(), 'No GPU found, please run without --cuda'
    torch.cuda.manual_seed_all(opt.seed)
    cudnn.benchmark = True



# model
model = SSD300(VOC.N_CLASSES)
cfg = model.config
#  Uncomment for test
#     model = torch.nn.DataParallel(model.cuda())


if opt.checkpoint:
    model.load_state_dict(torch.load(opt.checkpoint))
else:
    model.init_parameters(opt.backbone)

encoder = MultiBox(cfg)
criterion = MultiBoxLoss()

# cuda
if opt.cuda:
    model.cuda()

    criterion.cuda()
    cudnn.benchmark = True
    model = torch.nn.DataParallel(model, device_ids=[0,1])

# optimizer
optimizer = optim.SGD(model.parameters(), lr=opt.lr, momentum=0.9, weight_decay=5e-4)

# learning rate / iterations
init_lr = cfg.get('init_lr', 1e-3)
stepvalues = cfg.get('stepvalues', (80000, 100000))
# max_iter = cfg.get('max_iter', 120000)
max_iter = opt.max_iter

def adjust_learning_rate(optimizer, stage):
    lr = init_lr * (0.1 ** stage)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


# def warm_up(iteration):
# 	warmup_steps = [0, 300, 600, 1000]
# 	if iteration in warmup_steps:
# 		i = warmup_steps.index(iteration)
# 		lr = 10 ** np.linspace(-6, np.log10(init_lr), len(warmup_steps))[i]
# 		for param_group in optimizer.param_groups:
# 			param_group['lr'] = lr


def learning_rate_schedule(iteration):
    if iteration in stepvalues:
        adjust_learning_rate(optimizer, stepvalues.index(iteration) + 1)




def train():
    model.train()
    PRNG = RandomState(opt.seed)

    # augmentation / transforms
    transform = Compose([
            [ColorJitter(prob=0.5)],  # or write [ColorJitter(), None]
            BoxesToCoords(),
            Expand((1, 4), prob=0.5),
            ObjectRandomCrop(),
            HorizontalFlip(),
            Resize(300),
            CoordsToBoxes(),
            [SubtractMean(mean=VOC.MEAN)],
            [RGB2BGR()],
            [ToTensor()],
            ], PRNG, mode=None, fillval=VOC.MEAN)
    target_transform = encoder.encode

    dataset = VOCDetection(
            root=opt.voc_root,
            image_set=[('2007', 'trainval'), ('2012', 'trainval'),],
            keep_difficult=True,
            transform=transform,
            target_transform=target_transform)
    dataloader = DataLoader(dataset=dataset, batch_size=opt.batch_size, shuffle=True,
            num_workers=opt.threads, pin_memory=True)

    iteration = opt.start_iter
    while True:
        for input, loc, label in dataloader:
            t0 = time.time()

            learning_rate_schedule(iteration)

            input, loc, label = Variable(input), Variable(loc), Variable(label)
            if opt.cuda:
                input, loc, label = input.cuda(), loc.cuda(), label.cuda()

            xloc, xconf = model(input)
            loc_loss, conf_loss = criterion(xloc, xconf, loc, label)
            loss = loc_loss + conf_loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            t1 = time.time()

            if iteration % 10 == 0:
                summary.add_scalar('loss/loc_loss', loc_loss.data[0], iteration)
                summary.add_scalar('loss/conf_loss', conf_loss.data[0], iteration)
                summary.add_scalars('loss/loss', {"loc_loss": loc_loss.data[0],
                                                  "conf_loss": conf_loss.data[0],
                                                  "loss": loss.data[0]}, iteration)



            width_mult=opt.width_mult
            if iteration % 20 == 0 and iteration != opt.start_iter:
                print('iter: {}, Loss:{}, timer :{}'.format(iteration, loss.data[0], t1 - t0))

            if iteration % 5000 == 0 and iteration != opt.start_iter:
                print('Save state, iter: {}, Loss:{}'.format(iteration, loss.data[0]))
                if not os.path.isdir('weights'):
                        os.mkdir('weights')
                torch.save(model.state_dict(), 'weights/{}_width_mult_0.25_{}.pth'.format(cfg.get('name', 'SSD'), iteration))

            iteration += 1
            if iteration > max_iter:
                return 0





def test():
    PRNG = RandomState()

    dataset = VOCDetection(
            root=opt.voc_root,
            image_set=[('2007', 'test')],
            transform=Compose([
                BoxesToCoords(),
                Resize(300),
                CoordsToBoxes(),
                [SubtractMean(mean=VOC.MEAN)],
                [RGB2BGR()],
                [ToTensor()]]),
            target_transform=None)
    print(len(dataset))

    pred_bboxes, pred_labels, pred_scores, gt_bboxes, gt_labels = [], [], [], [], []
    for i in range(len(dataset)):
        num_images = len(dataset)
        t0 = time.time()


        img, loc, label = dataset[i]

        gt_bboxes.append(loc)
        gt_labels.append(label)

        input = Variable(img.unsqueeze(0), volatile=True)
        if opt.cuda:
            input = input.cuda()

        xloc, xconf = model(input)
        xloc = xloc.data.cpu().numpy()[0]
        xconf = xconf.data.cpu().numpy()[0]

        boxes, labels, scores = encoder.decode(xloc, xconf, nms_thresh=0.5, conf_thresh=0.01)

        pred_bboxes.append(boxes)
        pred_labels.append(labels)
        pred_scores.append(scores)
        t1 = time.time()

        print('Testing image {:d}/{:d}...., time:{}'.format(i+1, num_images,t1-t0))

    print(eval_voc_detection(pred_bboxes, pred_labels, pred_scores, gt_bboxes, gt_labels, iou_thresh=0.5, use_07_metric=True))



def demo():
    PRNG = RandomState()
    voc_vis = Viz()

    dataset = VOCDetection(
            root=opt.voc_root,
            image_set=[('2007', 'test')],
            transform=Compose([
                BoxesToCoords(),
                Resize(300),
                CoordsToBoxes(),
                [SubtractMean(mean=VOC.MEAN)],
                [RGB2BGR()],
                [ToTensor()]]),
            target_transform=None)
    print(len(dataset))

    i = PRNG.choice(len(dataset))
    for _ in range(1000):
        img, _, _ = dataset[i]

        input = Variable(img.unsqueeze(0), volatile=True)
        if opt.cuda:
            input = input.cuda()

        xloc, xconf = model(input)

        imgs = input.data.cpu().numpy().transpose(0, 2, 3, 1)
        xloc = xloc.data.cpu().numpy()
        xconf = xconf.data.cpu().numpy()

        for img, loc, conf in zip(imgs, xloc, xconf):
            #print(img.mean(axis=(0,1)))
            img = ((img[:, :, ::-1] + VOC.MEAN)).astype('uint8')
            boxes, labels, scores = encoder.decode(loc, conf, conf_thresh=0.5)

            img = voc_vis.draw_bbox(img, boxes, labels, True)
            cv2.imshow('0', img[:, :, ::-1])
            c = cv2.waitKey(0)
            if c == 27 or c == ord('q'):   # ESC / 'q'
                return
            else:
                i = PRNG.choice(len(dataset))





if __name__ == '__main__':
    if opt.test:
        test()
    elif opt.demo:
        demo()
    else:
        train()