typing for kwargs

torchmetrics/aggregation.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import warnings
-from typing import Any, Callable, Dict, List, Union
+from typing import Any, Callable, List, Union
 
 import torch
 from torch import Tensor
@@ -50,7 +50,7 @@ def __init__(
         fn: Union[Callable, str],
         default_value: Union[Tensor, List],
         nan_strategy: Union[str, float] = "error",
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         super().__init__(**kwargs)
         allowed_nan_strategy = ("error", "warn", "ignore")
@@ -122,7 +122,7 @@ class MaxMetric(BaseAggregator):
     def __init__(
         self,
         nan_strategy: Union[str, float] = "warn",
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         super().__init__(
             "max",
@@ -173,7 +173,7 @@ class MinMetric(BaseAggregator):
     def __init__(
         self,
         nan_strategy: Union[str, float] = "warn",
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         super().__init__(
             "min",
@@ -222,7 +222,7 @@ class SumMetric(BaseAggregator):
     def __init__(
         self,
         nan_strategy: Union[str, float] = "warn",
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         super().__init__(
             "sum",
@@ -271,7 +271,7 @@ class CatMetric(BaseAggregator):
     def __init__(
         self,
         nan_strategy: Union[str, float] = "warn",
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         super().__init__("cat", [], nan_strategy, **kwargs)
 
@@ -321,7 +321,7 @@ class MeanMetric(BaseAggregator):
     def __init__(
         self,
         nan_strategy: Union[str, float] = "warn",
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         super().__init__(
             "sum",

torchmetrics/audio/pesq.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict
+from typing import Any
 
 from torch import Tensor, tensor
 
@@ -80,7 +80,7 @@ def __init__(
         self,
         fs: int,
         mode: str,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         if not _PESQ_AVAILABLE:

torchmetrics/audio/pit.py

@@ -70,7 +70,7 @@ def __init__(
         self,
         metric_func: Callable,
         eval_func: str = "max",
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         base_kwargs: Dict[str, Any] = {
             "dist_sync_on_step": kwargs.pop("dist_sync_on_step", False),

torchmetrics/audio/sdr.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from torch import Tensor, tensor
 
@@ -88,7 +88,7 @@ def __init__(
         filter_length: int = 512,
         zero_mean: bool = False,
         load_diag: Optional[float] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
 
@@ -162,7 +162,7 @@ class ScaleInvariantSignalDistortionRatio(Metric):
     def __init__(
         self,
         zero_mean: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         self.zero_mean = zero_mean

torchmetrics/audio/snr.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict
+from typing import Any
 
 from torch import Tensor, tensor
 
@@ -69,7 +69,7 @@ class SignalNoiseRatio(Metric):
     def __init__(
         self,
         zero_mean: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         self.zero_mean = zero_mean
@@ -134,7 +134,7 @@ class ScaleInvariantSignalNoiseRatio(Metric):
 
     def __init__(
         self,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
 

torchmetrics/audio/stoi.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict
+from typing import Any
 
 from torch import Tensor, tensor
 
@@ -87,7 +87,7 @@ def __init__(
         self,
         fs: int,
         extended: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         if not _PYSTOI_AVAILABLE:

torchmetrics/classification/auc.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, List, Optional
+from typing import Any, List, Optional
 
 from torch import Tensor
 
@@ -44,7 +44,7 @@ class AUC(Metric):
     def __init__(
         self,
         reorder: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
 

torchmetrics/classification/auroc.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, List, Optional
+from typing import Any, List, Optional
 
 import torch
 from torch import Tensor
@@ -109,7 +109,7 @@ def __init__(
         pos_label: Optional[int] = None,
         average: Optional[str] = "macro",
         max_fpr: Optional[float] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
 

torchmetrics/classification/avg_precision.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, List, Optional, Union
 
 import torch
 from torch import Tensor
@@ -89,7 +89,7 @@ def __init__(
         num_classes: Optional[int] = None,
         pos_label: Optional[int] = None,
         average: Optional[str] = "macro",
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
 

torchmetrics/classification/binned_precision_recall.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -120,7 +120,7 @@ def __init__(
         self,
         num_classes: int,
         thresholds: Union[int, Tensor, List[float]] = 100,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
 
@@ -281,7 +281,7 @@ def __init__(
         num_classes: int,
         min_precision: float,
         thresholds: Union[int, Tensor, List[float]] = 100,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(num_classes=num_classes, thresholds=thresholds, **kwargs)
         self.min_precision = min_precision

torchmetrics/classification/calibration_error.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, List
+from typing import Any, List
 
 import torch
 from torch import Tensor
@@ -66,7 +66,7 @@ def __init__(
         self,
         n_bins: int = 15,
         norm: str = "l1",
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
 
         super().__init__(**kwargs)

torchmetrics/classification/cohen_kappa.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 import torch
 from torch import Tensor
@@ -77,7 +77,7 @@ def __init__(
         num_classes: int,
         weights: Optional[str] = None,
         threshold: float = 0.5,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         self.num_classes = num_classes

torchmetrics/classification/confusion_matrix.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 import torch
 from torch import Tensor
@@ -96,7 +96,7 @@ def __init__(
         normalize: Optional[str] = None,
         threshold: float = 0.5,
         multilabel: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         self.num_classes = num_classes

torchmetrics/classification/f_beta.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from torch import Tensor
 
@@ -260,7 +260,7 @@ def __init__(
         ignore_index: Optional[int] = None,
         top_k: Optional[int] = None,
         multiclass: Optional[bool] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(
             num_classes=num_classes,

torchmetrics/classification/hamming.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict
+from typing import Any
 
 import torch
 from torch import Tensor, tensor
@@ -65,7 +65,7 @@ class HammingDistance(Metric):
     def __init__(
         self,
         threshold: float = 0.5,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
 

torchmetrics/classification/hinge.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
 
 from torch import Tensor, tensor
 
@@ -97,7 +97,7 @@ def __init__(
         self,
         squared: bool = False,
         multiclass_mode: Optional[Union[str, MulticlassMode]] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
 

torchmetrics/classification/kl_divergence.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict
+from typing import Any
 
 import torch
 from torch import Tensor
@@ -74,7 +74,7 @@ def __init__(
         self,
         log_prob: bool = False,
         reduction: Literal["mean", "sum", "none", None] = "mean",
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         if not isinstance(log_prob, bool):

torchmetrics/classification/matthews_corrcoef.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict
+from typing import Any
 
 import torch
 from torch import Tensor
@@ -72,7 +72,7 @@ def __init__(
         self,
         num_classes: int,
         threshold: float = 0.5,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         self.num_classes = num_classes