rotary needs to be done with full precision to be safe

lucidrains · May 11, 2024 · 90be723 · 90be723
1 parent bca88e9
commit 90be723
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -6,7 +6,7 @@
 setup(
   name = 'vit-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '1.6.8',
+  version = '1.6.9',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   long_description=long_description,

diff --git a/vit_pytorch/rvt.py b/vit_pytorch/rvt.py
@@ -3,12 +3,14 @@
 import torch
 from torch import nn, einsum
 import torch.nn.functional as F
+from torch.cuda.amp import autocast
 
 from einops import rearrange, repeat
 from einops.layers.torch import Rearrange
 
 # rotary embeddings
 
+@autocast(enabled = False)
 def rotate_every_two(x):
     x = rearrange(x, '... (d j) -> ... d j', j = 2)
     x1, x2 = x.unbind(dim = -1)
@@ -22,6 +24,7 @@ def __init__(self, dim, max_freq = 10):
         scales = torch.linspace(1., max_freq / 2, self.dim // 4)
         self.register_buffer('scales', scales)
 
+    @autocast(enabled = False)
     def forward(self, x):
         device, dtype, n = x.device, x.dtype, int(sqrt(x.shape[-2]))