aljaz/DALLE2-pytorch
1
1
Fork 0
mirror of https://github.com/lucidrains/DALLE2-pytorch.git synced 2025-12-20 02:04:19 +01:00
Code Issues Packages Projects Releases Wiki Activity

remove einops exts for better pytorch 2.0 compile compatibility

Browse Source
This commit is contained in:
Phil Wang
2023-04-20 07:05:29 -07:00
parent 580274be79
commit 0069857cf8
4 changed files with 39 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
dalle2_pytorch/vqgan_vae.py
View File

@@ -11,8 +11,7 @@ import torch.nn.functional as F
from torch.autograd import grad as torch_grad
import torchvision
from einops import rearrange, reduce, repeat
from einops_exts import rearrange_many
from einops import rearrange, reduce, repeat, pack, unpack
from einops.layers.torch import Rearrange
# constants
@@ -408,7 +407,7 @@ class Attention(nn.Module):
x = self.norm(x)
q, k, v = self.to_qkv(x).chunk(3, dim = -1)
q, k, v = rearrange_many((q, k, v), 'b n (h d) -> b h n d', h = h)
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
q = q * self.scale
sim = einsum('b h i d, b h j d -> b h i j', q, k)