add typing for testing utils

tests/metrics/utils.py

@@ -1,9 +1,13 @@
-import torch
-import numpy as np
 import os
 import sys
 import pytest
 import pickle
+from typing import Callable
+
+import torch
+import numpy as np
+
+from pytorch_lightning.metrics import Metric
 
 NUM_PROCESSES = 2
 NUM_BATCHES = 10
@@ -17,7 +21,15 @@ def setup_ddp(rank, world_size):
     torch.distributed.init_process_group("gloo", rank=rank, world_size=world_size)
 
 
-def _compute_batch(rank, preds, target, metric_class, sk_metric, ddp_sync_on_step, worldsize=1, metric_args={}):
+def _compute_batch(rank: int,
+                   preds: torch.Tensor,
+                   target: torch.Tensor,
+                   metric_class: Metric,
+                   sk_metric: Callable,
+                   ddp_sync_on_step: bool,
+                   worldsize: int = 1,
+                   metric_args: dict = {}
+                   ):
     """ Utility function doing the actual comparison between lightning metric
         and reference metric
     """
@@ -56,7 +68,14 @@ def _compute_batch(rank, preds, target, metric_class, sk_metric, ddp_sync_on_ste
     assert np.allclose(result.numpy(), sk_result)
 
 
-def compute_batch(preds, target, metric_class, sk_metric, ddp_sync_on_step, ddp=False, metric_args={}):
+def compute_batch(preds: torch.Tensor,
+                  target: torch.Tensor,
+                  metric_class: Metric,
+                  sk_metric: Callable,
+                  ddp_sync_on_step: bool,
+                  ddp: bool = False,
+                  metric_args: dict = {}
+                  ):
     """ Utility function for comparing the result between a lightning class
         metric and another metric (often sklearns)