Merge pull request #7 from marksgraham/main

Import partial
lucidrains · Mar 5, 2023 · 9fa8d8e · 9fa8d8e
2 parents d041958 + ef4b809
commit 9fa8d8e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/memory_efficient_attention_pytorch/transformer.py b/memory_efficient_attention_pytorch/transformer.py
@@ -1,6 +1,7 @@
 import torch
 from torch import nn, einsum
 import torch.nn.functional as F
+from functools import partial
 
 from einops import rearrange
 from memory_efficient_attention_pytorch import FlashAttention, Attention