adding input assertions

ece.py

@@ -16,7 +16,7 @@ from typing import Dict
 
 import evaluate
 import datasets
-from torch import Tensor, LongTensor, amax
+from torch import Tensor, LongTensor
 from torchmetrics.functional.classification.calibration_error import (
     binary_calibration_error,
     multiclass_calibration_error,
@@ -109,14 +109,18 @@ class ECE(evaluate.Metric):
         references = LongTensor(references)
 
         # Determine number of classes / binary or multiclass
+        error_msg = "Expected to have predictions with shape (N,C,...) for multiclass or (N,...) for binary, " \
+                    f"and references with shape (N,...), but got {predictions.shape} and {references.shape}"
         binary = True
-        if predictions.dim() == references.dim() + 1:
+        if predictions.dim() == references.dim() + 1:  # multiclass
             binary = False
             if "num_classes" not in kwargs:
                 kwargs["num_classes"] = int(predictions.shape[1])
-        else:
-            raise ValueError("Bad input shape. Expected to have predictions with shape (N,C,...) and references"
-                             f"with shape (N,...), but got {predictions.shape} and {references.shape}")
+        elif predictions.dim() == references.dim() and "num_classes" in kwargs:
+            raise ValueError("You gave the num_classes argument, with predictions and references having the"
+                             "same number of dimensions. " + error_msg)
+        elif predictions.dim() != references.dim():
+            raise ValueError("Bad input shape. " + error_msg)
 
         # Compute the calibration
         if binary: