New train profile argument for loggers (#2862)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/yolo/utils/callbacks/neptune.py

@@ -1,9 +1,10 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
+
 import matplotlib.image as mpimg
 import matplotlib.pyplot as plt
 
 from ultralytics.yolo.utils import LOGGER, TESTS_RUNNING
-from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
+from ultralytics.yolo.utils.torch_utils import model_info_for_loggers
 
 try:
     import neptune
@@ -68,11 +69,7 @@ def on_train_epoch_end(trainer):
 def on_fit_epoch_end(trainer):
     """Callback function called at end of each fit (train+val) epoch."""
     if run and trainer.epoch == 0:
-        model_info = {
-            'parameters': get_num_params(trainer.model),
-            'GFLOPs': round(get_flops(trainer.model), 3),
-            'speed(ms)': round(trainer.validator.speed['inference'], 3)}
-        run['Configuration/Model'] = model_info
+        run['Configuration/Model'] = model_info_for_loggers(trainer)
     _log_scalars(trainer.metrics, trainer.epoch + 1)