Move curent_mask to perturbed tensor device

Summary: Currently `FeaturePermutation` and `FeatureAblation` both throw a device mismatch issue in https://fburl.com/code/9mfuidf4 because the `current_mask` is always created on CPU and never moved to the same device as `expanded_input` when CUDA is available.

Reviewed By: cyrjano

Differential Revision: D54969675

Author: Marcio Porto

captum/attr/_core/feature_ablation.py

@@ -559,6 +559,7 @@ def _construct_ablated_input(
         current_mask = torch.stack(
             [input_mask == j for j in range(start_feature, end_feature)], dim=0
         ).long()
+        current_mask = current_mask.to(expanded_input.device)
         ablated_tensor = (
             expanded_input * (1 - current_mask).to(expanded_input.dtype)
         ) + (baseline * current_mask.to(expanded_input.dtype))

captum/attr/_core/feature_permutation.py

@@ -301,6 +301,7 @@ def _construct_ablated_input(
         current_mask = torch.stack(
             [input_mask == j for j in range(start_feature, end_feature)], dim=0
         ).bool()
+        current_mask = current_mask.to(expanded_input.device)
 
         output = torch.stack(
             [