Merge pull request #3912 from MSciesiek/fix_casting_types_for_mps

Fix casting types for mps
fastai · May 26, 2023 · b4d8ac0 · b4d8ac0
2 parents 7176ffb + 8476639
commit b4d8ac0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/fastai/vision/augment.py b/fastai/vision/augment.py
@@ -557,7 +557,8 @@ def mask_tensor(
     if p==1.: return x
     if batch: return x if random.random() < p else x.new_zeros(*x.size()) + neutral
     if neutral != 0: x.add_(-neutral)
-    mask = x.new_empty(*x.size()).bernoulli_(p)
+    # Extra casting to float and long to prevent crashes on mps accelerator (issue #3911)
+    mask = x.new_empty(*x.size()).float().bernoulli_(p).long()
     x.mul_(mask)
     return x.add_(neutral) if neutral != 0 else x
 

diff --git a/nbs/09_vision.augment.ipynb b/nbs/09_vision.augment.ipynb
@@ -2073,7 +2073,8 @@
     "    if p==1.: return x\n",
     "    if batch: return x if random.random() < p else x.new_zeros(*x.size()) + neutral\n",
     "    if neutral != 0: x.add_(-neutral)\n",
-    "    mask = x.new_empty(*x.size()).bernoulli_(p)\n",
+    "    # Extra casting to float and long to prevent crashes on mps accelerator (issue #3911)\n",
+    "    mask = x.new_empty(*x.size()).float().bernoulli_(p).long()\n",
     "    x.mul_(mask)\n",
     "    return x.add_(neutral) if neutral != 0 else x"
    ]