test_pickle_accelerator should work now)

byi8220 · Sep 4, 2024 · c6a68c3 · c6a68c3
1 parent 1964011
commit c6a68c3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/accelerate/test_utils/scripts/test_distributed_data_loop.py b/src/accelerate/test_utils/scripts/test_distributed_data_loop.py
@@ -253,7 +253,8 @@ def test_pickle_accelerator():
     _ = accelerator.prepare(data_loader)
     pickled_accelerator = pickle.dumps(accelerator)
     unpickled_accelerator = pickle.loads(pickled_accelerator)
-    assert accelerator.state == unpickled_accelerator.state
+    # TODO: Maybe this should be implemented as __eq__ for AcceleratorState?
+    assert accelerator.state.__dict__ == unpickled_accelerator.state.__dict__
 
 
 def test_data_loader(data_loader, accelerator):