Update decoder.py

This commit is contained in:
jinzr 2023-09-20 01:34:49 +08:00
parent 7cc2dae940
commit c2adb8de3e

View File

@ -71,6 +71,10 @@ class Decoder(nn.Module):
groups=decoder_dim // 4, # group size == 4
bias=False,
)
else:
# To avoid `RuntimeError: Module 'Decoder' has no attribute 'conv'`
# when inference with torch.jit.script and context_size == 1
self.conv = nn.Identity()
def forward(self, y: torch.Tensor, need_pad: bool = True) -> torch.Tensor:
"""