Passing Model to CUDA

[mtd7]

Hello,

I've been working on trying to get your model to work as part of my Master's Dissertation and I've found some code that may no longer be up to date. net.cuda() didn't seem to be able to find my GPU; but I passed it with the code below and it seems to work...

net = make_model(n_classes=345, coord_input_dim=2, feat_input_dim=2, feat_dict_size=103, 
                 n_layers=network_configs['n_layers'], n_heads=network_configs['n_heads'], 
                 embed_dim=network_configs['embed_dim'], feedforward_dim=network_configs['feed_forward_hidden'], 
                 normalization=network_configs['normalization'], dropout=network_configs['dropout'], mlp_classifier_dropout=network_configs['mlp_classifier_dropout'])
#net = net.cuda()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = net.to(device)

This pattern in general needs to be repeated every where the .cuda() method is used...

Thanks.