Do #298 only for MPS (#308)

Lightning-AI · May 22, 2023 · 6e60162 · 6e60162
1 parent cdfbd91
commit 6e60162
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/generate.py b/generate.py
@@ -73,7 +73,8 @@ def generate(
             xm.mark_step()
 
         # concatenate the new generation
-        idx[t] = idx_next.item()
+        # https://github.com/pytorch/pytorch/issues/101936
+        idx[t] = idx_next.item() if idx.device.type == "mps" else idx_next
 
         # if <eos> token is triggered, return the output (stop generation)
         if idx_next == eos_id:

diff --git a/generate/full.py b/generate/full.py
@@ -61,7 +61,8 @@ def generate(
         idx_next = torch.multinomial(probs, num_samples=1)
 
         # concatenate the new generation
-        idx[t] = idx_next
+        # https://github.com/pytorch/pytorch/issues/101936
+        idx[t] = idx_next.item() if idx.device.type == "mps" else idx_next
 
         # if <eos> token is triggered, return the output (stop generation)
         if idx_next == eos_id: