Control for when output from model is a scalar or a 1D tensor

Summary: This is to make sure that we control for when the output is not a 2D tensor

Differential Revision: D69876980

Author: styusuf

captum/testing/helpers/basic_models.py

@@ -485,7 +485,14 @@ def forward(self, x: Tensor, add_input: Optional[Tensor] = None) -> Tensor:
 
         lin3_out = self.linear3(lin1_out_alt).to(torch.int64)
 
-        return torch.cat((lin2_out, lin3_out), dim=1)
+        output_tensors = torch.cat((lin2_out, lin3_out), dim=1)
+
+        # we return a dictionary of tensors as an output to test the case
+        # where an output accessor is required
+        return {
+            "task {}".format(i + 1): output_tensors[:, i]
+            for i in range(output_tensors.shape[1])
+        }
 
 
 class MultiRelu(nn.Module):