#conv0:
        for num_filters_pruned in range(16):
            model=prune_filter(model=model, layer=0, next_layer=3, indexes=index_pruned(model,0))
            if num_filters_pruned %4 == 0:
                model=model.cuda()
                for name, param in model.named_parameters():
                    device = param.device
                    if device == torch.device('cpu'):
                        print('ERROR in', name, device)
                train(model, train_loader, criterion, optimizer, device, 1)
                test(model,test_loader,criterion,device)

I added your piece of code but there is no “Error in” so I guess every parameter is in cuda