Ensure eval diff and mask tensor are on the same device (meta-pytorch#1542)

Summary: Pull Request resolved: meta-pytorch#1542 Failure (prod pkg): f718895429 Reviewed By: vivekmig Differential Revision: D72682579 fbshipit-source-id: 78045d6c58af10fae118204fcc4433d65d986eff

Author: sarahtranfb

captum/attr/_core/shapley_value.py

@@ -397,15 +397,19 @@ def attribute(
  if show_progress:
  attr_progress.update()
  if agg_output_mode:
- eval_diff = modified_eval - prev_results
+ eval_diff = (modified_eval - prev_results).to(
+ inputs_tuple[0].device
+ )
  prev_results = modified_eval
  else:
  # when perturb_per_eval > 1, every num_examples stands for
  # one perturb. Since the perturbs are from a consecutive
  # perumuation, each diff of a perturb is its eval minus
  # the eval of the previous perturb
  all_eval = torch.cat((prev_results, modified_eval), dim=0)
- eval_diff = all_eval[num_examples:] - all_eval[:-num_examples]
+ eval_diff = (
+ all_eval[num_examples:] - all_eval[:-num_examples]
+ ).to(inputs_tuple[0].device)
  prev_results = all_eval[-num_examples:]
 
  for j in range(len(total_attrib)):
@@ -689,7 +693,7 @@ def _evalFutToPrevResultsTuple(
  agg_output_mode,
  ) = prev_results_tuple
  if agg_output_mode:
- eval_diff = modified_eval - prev_results
+ eval_diff = (modified_eval - prev_results).to(inputs_tuple[0].device)
  prev_results = modified_eval
  else:
  # when perturb_per_eval > 1, every num_examples stands for
@@ -698,7 +702,9 @@ def _evalFutToPrevResultsTuple(
  # the eval of the previous perturb
 
  all_eval = torch.cat((prev_results, modified_eval), dim=0)
- eval_diff = all_eval[num_examples:] - all_eval[:-num_examples]
+ eval_diff = (all_eval[num_examples:] - all_eval[:-num_examples]).to(
+ inputs_tuple[0].device
+ )
  prev_results = all_eval[-num_examples:]
 
  for j in range(len(total_attrib)):
@@ -799,7 +805,10 @@ def _perturbation_generator(
  )
  current_tensors_list.append(current_tensors)
  current_mask_list.append(
- tuple(mask == feature_permutation[i] for mask in input_masks)
+ tuple(
+ (mask == feature_permutation[i]).to(inputs[0].device)
+ for mask in input_masks
+ )
  )
  if len(current_tensors_list) == perturbations_per_eval:
  combined_inputs = tuple(