Add type annotations in run_profiler.py

Author: ValerianRey

tests/profiling/run_profiler.py

@@ -3,7 +3,9 @@
 
 import torch
 from settings import DEVICE
+from torch import Tensor, nn
 from torch.profiler import ProfilerActivity, profile
+from torch.utils._pytree import PyTree
 from utils.architectures import (
     AlexNet,
     Cifar10Model,
@@ -105,15 +107,19 @@ def _save_and_print_trace(
 
 
 def profile_autojac(factory: ModuleFactory, batch_size: int) -> None:
-    def forward_backward_fn(model, inputs, loss_fn) -> None:
+    def forward_backward_fn(
+        model: nn.Module, inputs: PyTree, loss_fn: Callable[[PyTree], list[Tensor]]
+    ) -> None:
         aggregator = UPGrad()
         autojac_forward_backward(model, inputs, loss_fn, aggregator)
 
     profile_method("autojac", forward_backward_fn, factory, batch_size)
 
 
 def profile_autogram(factory: ModuleFactory, batch_size: int) -> None:
-    def forward_backward_fn(model, inputs, loss_fn) -> None:
+    def forward_backward_fn(
+        model: nn.Module, inputs: PyTree, loss_fn: Callable[[PyTree], list[Tensor]]
+    ) -> None:
         engine = Engine(model, batch_dim=0)
         weighting = UPGradWeighting()
         autogram_forward_backward(model, inputs, loss_fn, engine, weighting)