model.py

from torch import nn

from models import multimodal
import os

def generate_model(opt):


    
    assert opt.model in ['multimodal']

    if opt.model == 'multimodal':
        model = multimodal.MultiModal(opt.n_classes, fusion = opt.fusion, seq_length = opt.sample_duration, pretr_ef=opt.pretrain_path, num_heads=opt.num_heads)


    if opt.device != 'cpu':
        # model = model.to(opt.device)
        # model = nn.DataParallel(model, device_ids=None)
        
        # os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
        # os.environ["CUDA_VISIBLE_DEVICES"] = "1, 2"
        # device_ids = [0, 1]  
        # model = model.cuda()
        model = nn.DataParallel(model, device_ids=[0]).cuda()
        
        pytorch_total_params = sum(p.numel() for p in model.parameters() if
                               p.requires_grad)
        print("Total number of trainable parameters: ", pytorch_total_params)
        
    
    return model, model.parameters()