Commit 894d8925 authored by Phil Wang's avatar Phil Wang
Browse files

make sure dropout on post-softmax attention is present

parent 98723add

audiolm_pytorch/audiolm_pytorch.py

+3 −0
Original line number Diff line number Diff line
@@ -244,6 +244,8 @@ class Attention(nn.Module): 
        self.norm = LayerNorm(dim)

        self.context_norm = LayerNorm(dim_context) if norm_context else nn.Identity()



        self.attn_dropout = nn.Dropout(dropout)



        self.num_null_kv = num_null_kv

        self.null_kv = nn.Parameter(torch.randn(2, num_null_kv, dim_head))
@@ -298,6 +300,7 @@ class Attention(nn.Module): 
            sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)



        attn = sim.softmax(dim = -1)

        attn = self.attn_dropout(attn)



        out = einsum('b h i j, b j d -> b h i d', attn, v)

setup.py

+1 −1
Original line number Diff line number Diff line
@@ -3,7 +3,7 @@ from setuptools import setup, find_packages 
setup(

  name = 'audiolm-pytorch',

  packages = find_packages(exclude=[]),

  version = '0.1.7',

  version = '0.1.8',

  license='MIT',

  description = 'AudioLM - Language Modeling Approach to Audio Generation from Google Research - Pytorch',

  author = 'Phil Wang',