Commit d7f0153d authored by Anthony Larcher's avatar Anthony Larcher
Browse files

Merge branch 'dev_al' of https://git-lium.univ-lemans.fr/Larcher/sidekit into dev_al

parents c7f45d17 4414cb29
......@@ -1309,8 +1309,9 @@ def new_xtrain(dataset_description,
embedding_size = model.embedding_size
# Set the device and manage parallel processing
device = torch.cuda.device(local_rank)
torch.cuda.set_device(local_rank)
device = torch.device("cuda")
#device = torch.device("cuda")
model.to(device)
# If multi-gpu
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment