Hi @ptrblck, thanks for the reply. I have it already wrapped in train function and call it from the __main__ . I presume this is what you meant? Stil the same error. I tried everything from changing batch_size and reduced the number of parameters. But it throws an error before an epoch completes usually towards the end of epoch like below. also tried calling torch.cuda.empty_cache() inside the training loop with no success either.


Error Message : RuntimeError: CUDA out of memory. Tried to allocate 4.37 GiB (GPU 0; 11.17 GiB total capacity; 4.78 GiB already allocated; 1.58 GiB free; 4.49 GiB cached)===========
current memory allocated: 4693.4131
Max memory allocated: 9319.9683
Cached memory: 9490.0000
===========



def train(args):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    train_dataloader, test_dataloader, embedding_tuple = get_loader(args)
    model = SmallEmbedNN(embedding_tuple, len(args.cont_cols), 1,
                         emb_dropout=0.1, lin_layer_dropouts=0.1).to(device)
    
    criterion_l1 = L1Loss()
    
    # sparse optimizer for embeddings and normal optimizer for rest
    embed_sparse_vec = 'emb_layer.weight'
    opt = torch.optim.Adam([weights for name, weights in model.state_dict().items() \
                            if name!=embed_sparse_vec], lr=args.lr)
    optSparse = torch.optim.SparseAdam([model.state_dict()[embed_sparse_vec]], 
                                       lr=args.lr)
    train_loss_list=[]
    test_loss_list=[]
    
    try:
        for epoch in range(args.epochs):
            model.train()
            for _dict in tqdm(train_dataloader):
                # pass vars to cuda device
                _dict = {key:var.to(device) for key, var in _dict.items()}
                
                # zeroing parameter gradients
                opt.zero_grad()
                optSparse.zero_grad()

                # Forward Pass
                preds = model(_dict['cont'], _dict['cat'])
                loss_train_l1 = criterion_l1(preds, _dict['target'], 
                yr_weights=_dict['wts']) #L1 loss

                # Backward Pass and Optimization
                loss_train_l1.backward()
                opt.step()
                optSparse.step()

                # batch train loss
                train_loss = loss_train_l1.item()
                logger.info(f'Epoch:{epoch} train[batch] loss(L1):{train_loss:.4f}')
                # defrag cached memory
                torch.cuda.empty_cache()

            # append train loss
            train_loss_list.append(train_loss)

            # evaluate test-set and append test loss
            test_loss = eval_test(model, test_dataloader, criterion_l1, device)
            test_loss_list.append(test_loss)
            if epoch%1==0:
                logger.info(f'Epoch:{epoch}\ttrain_loss:{train_loss_list[-1]:.6f} \
                        \ttest_loss:{test_loss_list[-1]:.6f}')

            # earlystopping if test loss is not improving in last 10 epochs
            if min(test_loss_list) < min(test_loss_list[-20:]):
                break
    except Exception as e:
        logger.exception(e)
        memory_log(logger)