Skip to content

Commit ef4b809

Browse files
committed
Import partial
1 parent d041958 commit ef4b809

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

memory_efficient_attention_pytorch/transformer.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import torch
22
from torch import nn, einsum
33
import torch.nn.functional as F
4+
from functools import partial
45

56
from einops import rearrange
67
from memory_efficient_attention_pytorch import FlashAttention, Attention

0 commit comments

Comments
 (0)