Thanks for your reply!!
global training, prediction, cuda
training = True
prediction = True
cuda = True  
seed = 5

def run(opt="SGD", gpu=0, lr=0.001, lr_schedule=False, identifier='', scene_classification=False):

    # directories
    model_dir = r'.\model'
    train_dir = r'C:\Users\train_img'
    vali_dir = r'C:\Users\vali_img'
    test_dir = r'C:\Users\test_img'
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        print (torch.__version__)
        if not cuda:
            print("You have a CUDA device")
        else:
            torch.cuda.set_device(gpu)
            torch.cuda.manual_seed(seed)

    # building the net
    model = Unet_6(features=[32, 64]) 

    if cuda:
        model = model.cuda()
    new_identifier = '8' + identifier
    trainer = Trainer(net=model, train_dir=train_dir, vali_dir=vali_dir, test_dir=test_dir, model_dir=model_dir,
                      opt=opt, lr=lr, cuda=cuda, identifier=new_identifier, lr_schedule=lr_schedule)


    # training
    if training:
        bs = 8
        trainer.train_model(epoch=150, bs=bs)