parallel

nnet/xvector.py

@@ -1296,12 +1296,12 @@ def new_xtrain(dataset_description,
 
     # Display the entire configurations as YAML dictionnaries
     if local_rank < 1:
-    monitor.logger.info("\n*********************************\nDataset options\n*********************************\n")
-    monitor.logger.info(yaml.dump(dataset_opts, default_flow_style=False))
-    monitor.logger.info("\n*********************************\nModel options\n*********************************\n")
-    monitor.logger.info(yaml.dump(model_opts, default_flow_style=False))
-    monitor.logger.info("\n*********************************\nTraining options\n*********************************\n")
-    monitor.logger.info(yaml.dump(training_opts, default_flow_style=False))
+        monitor.logger.info("\n*********************************\nDataset options\n*********************************\n")
+        monitor.logger.info(yaml.dump(dataset_opts, default_flow_style=False))
+        monitor.logger.info("\n*********************************\nModel options\n*********************************\n")
+        monitor.logger.info(yaml.dump(model_opts, default_flow_style=False))
+        monitor.logger.info("\n*********************************\nTraining options\n*********************************\n")
+        monitor.logger.info(yaml.dump(training_opts, default_flow_style=False))
 
     # Initialize the model
     model = get_network(model_opts)
@@ -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