From 3c188e279135630a97927f3f55db328c15d8a6f3 Mon Sep 17 00:00:00 2001 From: Henry Isaacson Date: Mon, 19 Aug 2024 12:30:00 -0400 Subject: [PATCH] pytree --- src/beignet/func/__dataclass.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/beignet/func/__dataclass.py b/src/beignet/func/__dataclass.py index ee172d3f4e..ab2cd64ebf 100644 --- a/src/beignet/func/__dataclass.py +++ b/src/beignet/func/__dataclass.py @@ -1,7 +1,7 @@ import dataclasses from typing import List, Tuple, Type, TypeVar, Iterable -import optree +from torch.utils._pytree import register_pytree_node T = TypeVar("T") @@ -42,11 +42,10 @@ def _iterable_to_cls(meta, data) -> T: return dataclass_cls(**kwargs) - optree.register_pytree_node( + register_pytree_node( dataclass_cls, _iterate_cls, _iterable_to_cls, - namespace="beignet.func", ) return dataclass_cls