Update modeling_protonet.py

modeling_protonet.py

@@ -854,7 +854,7 @@ class AudioProtoNetPreTrainedModel(PreTrainedModel):
             nn.init.trunc_normal_(module.weight, std=0.02)
             if module.bias is not None:
                 nn.init.zeros_(module.bias)
-        if self.incorrect_class_connection is None and isinstance(self.last_layer, LinearLayerWithoutNegativeConnections): # TODO missing initilization
+        if isinstance(module, LinearLayerWithoutNegativeConnections) and self.incorrect_class_connection is None: # TODO missing initilization
             # Initialize all weights to the correct_class_connection value
             self.last_layer.weight.data.fill_(self.correct_class_connection)