From 4fc7084875ac672980c4e0207fa325c8c418a88e Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Wed, 1 Feb 2023 04:11:45 -0500 Subject: [PATCH] Fix a dimension bug in Transform2d (#2144) The dimension does not match when `inner_dim` is not equal to `in_channels`. --- src/diffusers/models/transformer_2d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/transformer_2d.py b/src/diffusers/models/transformer_2d.py index 57dd424aa4..2fdf70ce76 100644 --- a/src/diffusers/models/transformer_2d.py +++ b/src/diffusers/models/transformer_2d.py @@ -200,7 +200,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin): if self.is_input_continuous: # TODO: should use out_channels for continous projections if use_linear_projection: - self.proj_out = nn.Linear(in_channels, inner_dim) + self.proj_out = nn.Linear(inner_dim, in_channels) else: self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0) elif self.is_input_vectorized: