fix for loop

Author: NarineK

captum/influence/_core/tracincp.py

@@ -1060,7 +1060,7 @@ def get_checkpoint_contribution(checkpoint):
                                 torch.sum(
                                     layer_jacobian.flatten(start_dim=1) ** 2, dim=1
                                 )
-                                for layer_jacobian_flatten in layer_jacobians_flatten
+                                for layer_jacobian in layer_jacobians
                             ],
                             dim=0,
                         ),