Help using custom model with CEBRA

[Tsuchijo]

Hi, I am trying to use a custom model architecture with CEBRA but whenever I try to initialize the model I always get a similar error like:

     11 def __init__(self, num_neurons, num_units, num_output, normalize=True):
---> 12     super().__init__(
     13         nn.Conv1d(num_neurons, num_units, 2),
     14         nn.GELU(),
...
     17         num_input=num_neurons,
     18         num_output=num_output,
     19     )

TypeError: Model.__init__() takes 1 positional argument but 5 positional arguments (and 2 keyword-only arguments) were given

Which is odd im using the same formatting and model as in cebra.models. Here is the code registering the model:

import cebra.data.datatypes
import cebra.models.layers as cebra_layers
from cebra.models import register

@register("offset-5-model-test2")
class Offset1Model2dConv(cebra.models.Model, cebra.models.ConvolutionalModelMixin):
    """CEBRA model with a 5 sample receptive field and output normalization."""

    def __init__(self, num_neurons, num_units, num_output, normalize=True):
        super().__init__(
            nn.Conv1d(num_neurons, num_units, 2),
            nn.GELU(),
            cebra_layers._Skip(nn.Conv1d(num_units, num_units, 3), nn.GELU()),
            nn.Conv1d(num_units, num_output, 2),
            num_input=num_neurons,
            num_output=num_output,
        )

    def get_offset(self) -> cebra.data.datatypes.Offset:
        """See :py:meth:`~.Model.get_offset`"""
        return cebra.data.Offset(2, 3)

Any help would be appreciated, thanks!

[Tsuchijo]

Ended up figuring out the issue, I previously registered a model with the same class name that was incorrectly formatted which was causing an error.