make mypy happy

Author: ieee8023

captum/attr/_core/latent_shift.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-from typing import Callable, Tuple
+from typing import Callable, Tuple, List
 
 import os
 import shutil
@@ -53,7 +53,7 @@ def __init__(self, forward_func: Callable, autoencoder) -> None:
     def attribute(
         self,
         inputs: Tensor,
-        target: int = None,
+        target: int,
         fix_range: Tuple = None,
         search_pred_diff: float = 0.8,
         search_step_size: float = 10.0,
@@ -197,21 +197,20 @@ def compute_shift(lambdax):
             lbound,
             rbound,
             np.abs((lbound - rbound) / lambda_sweep_steps)
-        )
+        ).tolist()
 
         preds = []
         generated_images = []
 
         for lam in lambdas:
             x_lambdax, pred = compute_shift(lam)
             generated_images.append(x_lambdax.cpu().numpy())
-            preds.append(pred)
+            preds.append(float(pred))
 
         params = {}
         params['generated_images'] = generated_images
         params['lambdas'] = lambdas
         params['preds'] = preds
-        params['target'] = target
 
         x_lambda0 = x_lambda0.detach().cpu().numpy()
         if heatmap_method == 'max':
@@ -276,9 +275,6 @@ def generate_video(
                 return a video to show in a jupyter notebook.
         """
 
-        if not target_filename:
-            target_filename = f'video-{params["target"]}'
-
         if os.path.exists(target_filename + ".mp4"):
             os.remove(target_filename + ".mp4")