Update celle/celle.py

celle/celle.py

@@ -262,8 +262,10 @@ class ModelExtender(nn.Module):
 
         # Set the number of output features and initialize the scaling layer
         self.out_features = out_features
-        self.scale_layer = nn.Linear(self.in_features, self.out_features)
-
+        if self.in_features != self.out_features:
+            self.scale_layer = nn.Linear(self.in_features, self.out_features)
+        else:
+            self.scale_layer = nn.Identity()
         # Determine whether to freeze the model's parameters
         self.fixed_embedding = fixed_embedding
         if self.fixed_embedding: