From 57c5d02bbed8ad16b4ac3f8903d106e978448431 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 29 Jun 2021 16:03:10 +0200 Subject: [PATCH] Concise `TransformerBlock()` (#3821) --- models/common.py | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/models/common.py b/models/common.py index 4211db4..96d63a0 100644 --- a/models/common.py +++ b/models/common.py @@ -77,18 +77,8 @@ class TransformerBlock(nn.Module): if self.conv is not None: x = self.conv(x) b, _, w, h = x.shape - p = x.flatten(2) - p = p.unsqueeze(0) - p = p.transpose(0, 3) - p = p.squeeze(3) - e = self.linear(p) - x = p + e - - x = self.tr(x) - x = x.unsqueeze(3) - x = x.transpose(0, 3) - x = x.reshape(b, self.c2, w, h) - return x + p = x.flatten(2).unsqueeze(0).transpose(0, 3).squeeze(3) + return self.tr(p + self.linear(p)).unsqueeze(3).transpose(0, 3).reshape(b, self.c2, w, h) class Bottleneck(nn.Module):