mps fix (#1227)

Jerry-Master · facebook-github-bot · commit 7ee195d4d1e6 · 2024-02-22T07:28:07.000-08:00
Summary: By adding that single conversion from float64 to float32 the integrated gradients is fully compatible with MPS backend, which it wasn't previously. The change can be accepted as valid for any backend since float64 is a precision almost nobody using nowadays in deep learning. It seems that the default for torch.tensor is not updated with current trends. Pull Request resolved: #1227 Reviewed By: cyrjano Differential Revision: D54047349 Pulled By: vivekmig fbshipit-source-id: 1ffda83f065a3f14fa3c5b0229fe0feb5035cc99
diff --git a/captum/attr/_core/integrated_gradients.py b/captum/attr/_core/integrated_gradients.py
@@ -359,7 +359,7 @@ def _attribute(
         # calling contiguous to avoid `memory whole` problems
         scaled_grads = [
             grad.contiguous().view(n_steps, -1)
-            * torch.tensor(step_sizes).view(n_steps, 1).to(grad.device)
+            * torch.tensor(step_sizes).float().view(n_steps, 1).to(grad.device)
             for grad in grads
         ]
 

Original file line number	Diff line number	Diff line change
`@@ -359,7 +359,7 @@ def _attribute(`
`359`	`359`	# calling contiguous to avoid `memory whole` problems
`360`	`360`	`scaled_grads = [`
`361`	`361`	`grad.contiguous().view(n_steps, -1)`
`362`		`- * torch.tensor(step_sizes).view(n_steps, 1).to(grad.device)`
	`362`	`+ * torch.tensor(step_sizes).float().view(n_steps, 1).to(grad.device)`
`363`	`363`	`for grad in grads`
`364`	`364`	`]`
`365`	`365`