From 376a934ced841256db41eebdebedeb827439331a Mon Sep 17 00:00:00 2001
From: Eran <exx8eran@gmail.com>
Date: Sat, 7 Jan 2023 00:40:22 +0200
Subject: [PATCH] Update train_cls.py

add parallelism to the classification training.
---
 train_cls.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/train_cls.py b/train_cls.py
index 22e8200..e74e446 100644
--- a/train_cls.py
+++ b/train_cls.py
@@ -17,7 +17,7 @@
 import shutil
 import hydra
 import omegaconf
-
+import torch
 
 def test(model, loader, num_class=40):
     mean_correct = []
@@ -66,6 +66,7 @@ def main(args):
     shutil.copy(hydra.utils.to_absolute_path('models/{}/model.py'.format(args.model.name)), '.')
 
     classifier = getattr(importlib.import_module('models.{}.model'.format(args.model.name)), 'PointTransformerCls')(args).cuda()
+    classifier = torch.nn.DataParallel(classifier)
     criterion = torch.nn.CrossEntropyLoss()
 
     try:
@@ -159,4 +160,4 @@ def main(args):
     logger.info('End of training...')
 
 if __name__ == '__main__':
-    main()
\ No newline at end of file
+    main()