diff --git a/src/zenml/client.py b/src/zenml/client.py index 0441db7b973..b3dddd3e1c9 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -109,7 +109,6 @@ EventSourceResponse, EventSourceUpdate, FlavorFilter, - FlavorRequest, FlavorResponse, ModelFilter, ModelRequest, @@ -2201,17 +2200,8 @@ def create_flavor( "configuration class' docstring." ) - create_flavor_request = FlavorRequest( - source=source, - type=flavor.type, - name=flavor.name, - config_schema=flavor.config_schema, - integration="custom", - user=self.active_user.id, - workspace=self.active_workspace.id, - ) - - return self.zen_store.create_flavor(flavor=create_flavor_request) + flavor_request = flavor.to_model(integration="custom", is_custom=True) + return self.zen_store.create_flavor(flavor=flavor_request) def get_flavor( self,