diff --git a/torch_uncertainty/utils/evaluation_loop.py b/torch_uncertainty/utils/evaluation_loop.py index d079cae4..f3088103 100644 --- a/torch_uncertainty/utils/evaluation_loop.py +++ b/torch_uncertainty/utils/evaluation_loop.py @@ -27,7 +27,7 @@ def _add_row(table: Table, metric_name: str, value: Tensor) -> None: if metric_name in PERCENTAGE_METRICS: value = value * 100 - table.add_row(metric_name, f"{value.item():.2f}%") + table.add_row(metric_name, f"{value.item():.3f}%") else: table.add_row(metric_name, f"{value.item():.5f}") diff --git a/torch_uncertainty/utils/to_hub_format.py b/torch_uncertainty/utils/to_hub_format.py index 1e1136d3..8db9bcea 100644 --- a/torch_uncertainty/utils/to_hub_format.py +++ b/torch_uncertainty/utils/to_hub_format.py @@ -24,7 +24,7 @@ raise ValueError("File does not exist") dtype = torch.float16 if args.fp16 else torch.float32 -model = torch.load(args.path)["state_dict"] +model = torch.load(args.path, weights_only=True)["state_dict"] model = {key.replace("model.", ""): val.to(device="cpu", dtype=dtype) for key, val in model.items()} output_name = args.name