Fix mypy type errors in generate_paper_figures and train_kd

Author: dariocazzani

analysis/generate_paper_figures.py

@@ -229,7 +229,7 @@ def figure_gap_scaling(output: Path) -> None:
 
     # Custom x-axis
     ax.set_xticks(log_classes)
-    ax.set_xticklabels(classes)
+    ax.set_xticklabels([str(c) for c in classes])
     ax.set_xlim(0.7, 2.6)
     ax.set_ylim(0, 10)
 

experiments/train_kd.py

@@ -104,7 +104,8 @@ def train_kd(config: TrainConfig, teacher_path: Path, temperature: float, alpha:
     # Skip if already completed
     if results_path.exists():
         log.warning("Skipping %s (already completed)", output_dir)
-        return json.loads(results_path.read_text())
+        result: dict = json.loads(results_path.read_text())
+        return result
 
     set_seed(config.seed)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")