Fix type issues

Author: vivekmig

captum/insights/attr_vis/server.py

@@ -6,6 +6,8 @@
 from time import sleep
 from typing import cast, Dict, Optional
 
+from captum.insights import AttributionVisualizer
+
 from captum.log import log_usage
 from flask import Flask, jsonify, render_template, request
 from flask.wrappers import Response
@@ -44,7 +46,7 @@ def attribute() -> Response:
     r = cast(Dict, request.get_json(force=True))
     return jsonify(
         namedtuple_to_dict(
-            visualizer._calculate_attribution_from_cache(  # type: ignore
+            cast(AttributionVisualizer, visualizer)._calculate_attribution_from_cache(
                 r["inputIndex"], r["modelIndex"], r["labelIndex"]
             )
         )
@@ -54,15 +56,15 @@ def attribute() -> Response:
 @app.route("/fetch", methods=["POST"])
 def fetch() -> Response:
     # force=True needed, see comment for "/attribute" route above
-    visualizer._update_config(request.get_json(force=True))  # type: ignore
-    visualizer_output = visualizer.visualize()  # type: ignore
+    cast(AttributionVisualizer, visualizer)._update_config(request.get_json(force=True))
+    visualizer_output = cast(AttributionVisualizer, visualizer).visualize()
     clean_output = namedtuple_to_dict(visualizer_output)
     return jsonify(clean_output)
 
 
 @app.route("/init")
 def init() -> Response:
-    return jsonify(visualizer.get_insights_config())  # type: ignore
+    return jsonify(cast(AttributionVisualizer, visualizer).get_insights_config())
 
 
 @app.route("/")

scripts/install_via_pip.sh

@@ -12,7 +12,7 @@ while getopts 'ndfv:' flag; do
     d) DEPLOY=true ;;
     f) FRAMEWORKS=true ;;
     v) CHOSEN_TORCH_VERSION=${OPTARG};;
-    *) echo "usage: $0 [-n] [-d] [-f] [-v version] [-m install_mode]" >&2
+    *) echo "usage: $0 [-n] [-d] [-f] [-v version]" >&2
        exit 1 ;;
     esac
   done

tests/helpers/basic_models.py

@@ -43,8 +43,8 @@ class BasicModel(nn.Module):
     def __init__(self) -> None:
         super().__init__()
 
-    def forward(self, input: int):
-        input = 1 - F.relu(torch.tensor(1 - input))
+    def forward(self, input: Tensor):
+        input = 1 - F.relu(1 - input)
         return input