fix linting

pytorch · Feb 13, 2024 · a2d5c08 · a2d5c08
1 parent 9a7e7f4
commit a2d5c08
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/torchtrain/metrics_utils.py b/torchtrain/metrics_utils.py
@@ -1,14 +1,13 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 
-import torch
 import torch.nn as nn
 
 def get_num_params(model: nn.Module, only_trainable: bool = False)-> int:
-    """ Get the total model params
-    args: only_trainable - only count trainable params
     """
-
+    Get the total model params
+    Args: only_trainable: whether to only count trainable params
+    """
     param_list = list(model.parameters())
     if only_trainable:
         param_list = [p for p in param_list if p.requires_grad]