def METHOD_NAME(mx_model):
    class CustomModel(Model):
        def fn(self) -> int:
            return 1337
    model = MXNetWrapper(mx_model, model_class=CustomModel)
    assert isinstance(model, CustomModel)
    assert model.fn() == 1337