Update inspect_gen.py

OverLordGoldDragon · May 30, 2020 · ce34d0e · ce34d0e
1 parent 2f699e7
commit ce34d0e
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/see_rnn/inspect_gen.py b/see_rnn/inspect_gen.py
@@ -159,7 +159,12 @@ def _get_info(model, _id, layer, mode):
     else:
         grads_fn = _make_grads_fn(model, params=params)
         if sample_weight is None:
-            sw = None if TF_KERAS else np.ones(len(input_data))
+            if TF_KERAS:
+                sw = None
+            elif isinstance(input_data, list):
+                sw = np.ones(len(input_data[0]))
+            else:
+                sw = np.ones(len(input_data))
             if isinstance(input_data, list):
                 sample_weight = []
                 for x in input_data: