Ablation Device Fix (#528)

Summary:
Pull Request resolved: #528

Allows output to be on a different device than input by moving output difference to input device.

Reviewed By: miguelmartin75

Differential Revision: D25001273

fbshipit-source-id: a9b6d8e8bb585d5360c53272a5502f4e8f257459

Author: vivekmig

captum/attr/_core/feature_ablation.py

@@ -346,6 +346,7 @@ def attribute(
                         eval_diff = (
                             initial_eval - modified_eval.reshape((-1, num_outputs))
                         ).reshape((-1, num_outputs) + (len(inputs[i].shape) - 1) * (1,))
+                        eval_diff = eval_diff.to(total_attrib[i].device)
                     if self.use_weights:
                         weights[i] += current_mask.float().sum(dim=0)
                     total_attrib[i] += (eval_diff * current_mask.to(attrib_type)).sum(