diff --git a/src/nnsight/schema/Request.py b/src/nnsight/schema/Request.py index 5d7adaab..4fa6e317 100644 --- a/src/nnsight/schema/Request.py +++ b/src/nnsight/schema/Request.py @@ -21,7 +21,7 @@ class RequestModel(BaseModel): arbitrary_types_allowed=True, protected_namespaces=() ) - object: str | OBJECT_TYPES + object: Union[str, OBJECT_TYPES] model_key: str id: str = None diff --git a/src/nnsight/schema/format/types.py b/src/nnsight/schema/format/types.py index be97f25b..245bf889 100644 --- a/src/nnsight/schema/format/types.py +++ b/src/nnsight/schema/format/types.py @@ -77,8 +77,8 @@ def deserialize(self, handler: DeserializeHandler) -> Node: target: Union[FunctionModel, FunctionType] args: List[ValueTypes] = [] kwargs: Dict[str, ValueTypes] = {} - condition: None | Union[ - NodeReferenceType, NodeModel.Reference + condition: Union[ + NodeReferenceType, NodeModel.Reference, None ] = None @model_serializer(mode='wrap')