Merge pull request #276 from roboflow/fix-upload-loaded-model

Save uploaded num classes

Author: probicheaux

rfdetr/main.py

@@ -164,6 +164,10 @@ def train(self, callbacks: DefaultDict[str, List[Callable]], **kwargs):
                     f"Currently supported callbacks: {currently_supported_callbacks}"
                 )
         args = populate_args(**kwargs)
+        if getattr(args, 'class_names') is not None:
+            self.args.class_names = args.class_names
+            self.args.num_classes = args.num_classes
+
         utils.init_distributed_mode(args)
         print("git:\n  {}\n".format(utils.get_sha()))
         print(args)