def __setstate__(self, state_dict):
for k, v in state_dict.items():
if isinstance(v, tuple) and len(v) == 2 and isinstance(
v[1], SerializationType):
self.__dict__[k] = deserialize_fn(v)
else:
self.__dict__[k] = v
def __setstate__(self, state):
(
self.main_datapipe,
self.num_instances,
self.buffer_size,
serialized_fn_with_method,
self.drop_none,
) = state
self.classifier_fn = deserialize_fn(serialized_fn_with_method)
self._datapipe_iterator = None
self.current_buffer_usage = 0
self.child_buffers = [deque() for _ in range(self.num_instances)]
self.instance_started = [False] * self.num_instances
self.main_datapipe_exhausted = False
def __setstate__(self, state):
(
self.main_datapipe,
self.num_instances,
self.buffer_size,
fn,
self.drop_none,
) = state
if isinstance(fn, tuple) and len(fn) == 2 and isinstance(
fn[1], SerializationType):
self.classifier_fn = deserialize_fn(fn)
else:
self.classifier_fn = fn
self._datapipe_iterator = None
self.current_buffer_usage = 0
self.child_buffers = [deque() for _ in range(self.num_instances)]
self.instance_started = [False] * self.num_instances
self.main_datapipe_exhausted = False