From c9472ed3f9df5a3419e442bc03f093a3bd9f2940 Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Fri, 20 Dec 2024 00:48:11 -0800 Subject: [PATCH 01/33] unify `get_calculator` function --- mlip_arena/models/utils.py | 10 +- mlip_arena/tasks/eos_alloy/flow.py | 3 +- mlip_arena/tasks/eos_alloy/run.ipynb | 53921 ++++++++++++++++++++++++- mlip_arena/tasks/md.py | 80 +- mlip_arena/tasks/optimize.py | 44 +- mlip_arena/tasks/utils.py | 47 + tests/test_eos.py | 2 +- 7 files changed, 53993 insertions(+), 114 deletions(-) create mode 100644 mlip_arena/tasks/utils.py diff --git a/mlip_arena/models/utils.py b/mlip_arena/models/utils.py index 8a2c697..ce4cb19 100644 --- a/mlip_arena/models/utils.py +++ b/mlip_arena/models/utils.py @@ -5,12 +5,12 @@ def get_freer_device() -> torch.device: """Get the GPU with the most free memory, or use MPS if available. - s - Returns: - torch.device: The selected GPU device or MPS. - Raises: - ValueError: If no GPU or MPS is available. + Returns: + torch.device: The selected GPU device or MPS. + + Raises: + ValueError: If no GPU or MPS is available. """ device_count = torch.cuda.device_count() if device_count > 0: diff --git a/mlip_arena/tasks/eos_alloy/flow.py b/mlip_arena/tasks/eos_alloy/flow.py index c516c66..dc2bdb1 100644 --- a/mlip_arena/tasks/eos_alloy/flow.py +++ b/mlip_arena/tasks/eos_alloy/flow.py @@ -72,8 +72,7 @@ def save_to_hdf( family_path = Path(__file__).parent / REGISTRY[calculator_name]["family"] family_path.mkdir(parents=True, exist_ok=True) - with open(family_path / f"{calculator_name}_{formula}.json", "w") as f: - json.dump(result, f, indent=2) + df.to_json(family_path / f"{calculator_name}_{formula}.json", indent=2) with SafeHDFStore(fpath, mode="a") as store: store.append( diff --git a/mlip_arena/tasks/eos_alloy/run.ipynb b/mlip_arena/tasks/eos_alloy/run.ipynb index 217b1fb..cb5571f 100644 --- a/mlip_arena/tasks/eos_alloy/run.ipynb +++ b/mlip_arena/tasks/eos_alloy/run.ipynb @@ -85,7 +85,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "metadata": { "tags": [] }, @@ -98,7 +98,7 @@ "\n", "#SBATCH -A m3828\n", "#SBATCH --mem=0\n", - "#SBATCH -t 00:30:00\n", + "#SBATCH -t 01:00:00\n", "#SBATCH -J eos\n", "#SBATCH -q regular\n", "#SBATCH -N 1\n", @@ -107,18 +107,18 @@ "source ~/.bashrc\n", "module load python\n", "source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\n", - "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.37:38383 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 60\n", + "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.21:41405 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 60\n", "\n" ] }, { "data": { "text/html": [ - "
22:51:52.414 | INFO    | prefect.engine - Created flow run 'sly-owl' for flow 'run-from-db'\n",
+       "
22:23:21.156 | INFO    | prefect.engine - Created flow run 'benign-parrot' for flow 'run-from-db'\n",
        "
\n"
       ],
       "text/plain": [
-       "22:51:52.414 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'sly-owl'\u001b[0m for flow\u001b[1;35m 'run-from-db'\u001b[0m\n"
+       "22:23:21.156 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'benign-parrot'\u001b[0m for flow\u001b[1;35m 'run-from-db'\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -127,11 +127,11 @@
     {
      "data": {
       "text/html": [
-       "
22:51:52.419 | INFO    | prefect.engine - View at https://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/runs/flow-run/d2afea80-afd7-4be2-b734-087f31274217\n",
+       "
22:23:21.160 | INFO    | prefect.engine - View at https://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/runs/flow-run/a449ef94-1e7d-4bf6-940c-36e829292a73\n",
        "
\n"
       ],
       "text/plain": [
-       "22:51:52.419 | \u001b[36mINFO\u001b[0m    | prefect.engine - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/runs/flow-run/d2afea80-afd7-4be2-b734-087f31274217\u001b[0m\n"
+       "22:23:21.160 | \u001b[36mINFO\u001b[0m    | prefect.engine - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/runs/flow-run/a449ef94-1e7d-4bf6-940c-36e829292a73\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -140,11 +140,11 @@
     {
      "data": {
       "text/html": [
-       "
22:51:53.100 | INFO    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(f5a25e18, 'tcp://128.55.64.37:38383', workers=0, threads=0, memory=0 B)\n",
+       "
22:23:21.802 | INFO    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(28696f19, 'tcp://128.55.64.21:41405', workers=0, threads=0, memory=0 B)\n",
        "
\n"
       ],
       "text/plain": [
-       "22:51:53.100 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(f5a25e18, 'tcp://128.55.64.37:38383', workers=0, threads=0, memory=0 B)\n"
+       "22:23:21.802 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(28696f19, 'tcp://128.55.64.21:41405', workers=0, threads=0, memory=0 B)\n"
       ]
      },
      "metadata": {},
@@ -153,11 +153,11 @@
     {
      "data": {
       "text/html": [
-       "
22:52:10.768 | INFO    | Task run 'get_atoms_from_db-f27' - Created task run 'get_atoms_from_db-f27' for task 'get_atoms_from_db'\n",
+       "
22:23:43.657 | INFO    | Task run 'get_atoms_from_db-7a4' - Created task run 'get_atoms_from_db-7a4' for task 'get_atoms_from_db'\n",
        "
\n"
       ],
       "text/plain": [
-       "22:52:10.768 | \u001b[36mINFO\u001b[0m    | Task run 'get_atoms_from_db-f27' - Created task run 'get_atoms_from_db-f27' for task 'get_atoms_from_db'\n"
+       "22:23:43.657 | \u001b[36mINFO\u001b[0m    | Task run 'get_atoms_from_db-7a4' - Created task run 'get_atoms_from_db-7a4' for task 'get_atoms_from_db'\n"
       ]
      },
      "metadata": {},
@@ -166,15 +166,53900 @@
     {
      "data": {
       "text/html": [
-       "
22:52:41.782 | INFO    | Task run 'get_atoms_from_db-f27' - Finished in state Completed()\n",
+       "
22:24:15.703 | INFO    | Task run 'get_atoms_from_db-7a4' - Finished in state Completed()\n",
        "
\n"
       ],
       "text/plain": [
-       "22:52:41.782 | \u001b[36mINFO\u001b[0m    | Task run 'get_atoms_from_db-f27' - Finished in state \u001b[32mCompleted\u001b[0m()\n"
+       "22:24:15.703 | \u001b[36mINFO\u001b[0m    | Task run 'get_atoms_from_db-7a4' - Finished in state \u001b[32mCompleted\u001b[0m()\n"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2024-12-19 22:53:36,546 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:03:03,741 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:23:02,540 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:23:02,653 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:23:03,677 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:23:04,051 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:23:04,152 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:23:04,575 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:23:55,288 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:23:55,468 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:23:56,465 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:26:36,531 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:26:36,892 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:26:37,872 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:26:37,956 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:26:38,324 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:26:39,665 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:26:39,908 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:26:41,066 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:26:42,023 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:26:43,022 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:26:44,439 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:26:45,410 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:26:45,791 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:26:46,825 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:26:47,177 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:26:48,034 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:26:48,406 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:26:48,871 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:28:53,763 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:28:55,986 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:28:56,948 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:28:58,018 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:28:58,980 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:00,210 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:00,621 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:01,081 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:01,164 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:01,580 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:34,753 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:35,741 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:36,085 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:37,369 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:38,339 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:38,718 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:40,669 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:41,598 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:42,601 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:44,035 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:44,944 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:45,925 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:46,103 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:46,752 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:47,193 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:48,875 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:49,825 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:50,838 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:52,322 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:29:54,548 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:29:58,236 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:00,650 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:02,596 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:03,544 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:04,537 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:05,272 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:06,194 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:07,186 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:07,652 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:09,518 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:09,611 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:10,575 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:10,663 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:11,268 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:14,435 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:15,359 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:16,438 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:16,546 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:16,928 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:17,976 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:20,384 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:23,110 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:23,222 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:24,387 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:26,722 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:26,827 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:27,464 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:29,162 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:29,714 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:30,848 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:31,234 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:31,608 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:32,593 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:35,430 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:38,186 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:39,336 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:39,935 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:40,304 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:41,341 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:43,157 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:44,941 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:45,937 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:46,573 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:46,704 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:48,933 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:49,893 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:51,825 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:53,097 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:54,073 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:56,347 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:30:56,434 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:56,985 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:57,861 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:30:59,078 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:00,027 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:00,994 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:01,077 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:01,156 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:02,104 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:02,652 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:04,827 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:04,989 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:06,302 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:07,052 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:09,555 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:10,511 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:12,127 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:12,744 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:12,915 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:13,099 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:14,082 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:14,742 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:16,081 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:16,170 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:17,145 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:35,797 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:36,350 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:37,323 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:39,772 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:40,001 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:40,368 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:42,326 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:42,873 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:42,961 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:43,046 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:43,950 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:45,099 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:46,050 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:46,593 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:49,527 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:49,880 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:50,257 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:50,348 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:50,704 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:52,845 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:52,930 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:53,297 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:53,687 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:54,067 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:54,445 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:55,387 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:55,748 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:55,840 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:56,786 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:57,712 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:59,069 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:31:59,447 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:31:59,540 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:00,452 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:00,817 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:01,204 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:02,162 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:02,540 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:02,936 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:03,873 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:04,245 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:04,342 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:04,898 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:04,986 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:05,378 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:06,793 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:06,879 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:07,247 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:07,648 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:08,057 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:08,460 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:08,559 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:08,646 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:09,027 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:09,420 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:10,817 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:11,362 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:11,600 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:12,549 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:12,927 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:13,310 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:14,303 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:14,685 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:15,628 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:16,575 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:17,514 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:17,888 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:17,984 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:18,922 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:19,297 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:19,694 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:19,795 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:20,183 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:21,294 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:21,664 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:21,761 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:22,130 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:22,227 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:23,094 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:24,043 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:24,407 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:25,364 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:25,748 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:26,672 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:27,059 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:27,459 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:28,381 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:28,468 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:28,859 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:29,792 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:31,214 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:32,296 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:32,380 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:32,466 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:33,837 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:34,752 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:35,139 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:36,045 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:36,583 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:36,967 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:37,061 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:37,142 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:38,078 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:39,008 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:39,840 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:40,222 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:41,181 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:41,552 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:41,950 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:42,044 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:42,981 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:43,516 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:43,612 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:43,693 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:44,063 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:44,446 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:44,541 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:45,080 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:45,485 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:45,892 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:46,831 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:47,215 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:47,311 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:48,135 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:49,070 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:49,603 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:50,570 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:50,652 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:51,599 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:32:51,679 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:52,477 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:52,852 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:53,207 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:32:58,955 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:01,072 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:01,431 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:02,334 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:02,695 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:03,242 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:03,599 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:04,392 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:04,922 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:05,271 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:06,639 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:06,726 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:07,625 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:07,705 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:08,642 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:08,730 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:09,687 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:10,061 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:11,369 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:11,453 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:11,533 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:12,452 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:13,405 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:14,306 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:15,257 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:16,200 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:16,578 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:17,480 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:17,563 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:17,652 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:17,921 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:18,008 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:18,391 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:18,480 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:18,857 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:19,241 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:19,631 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:20,016 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:20,108 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:20,192 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:21,131 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:21,401 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:21,772 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:22,160 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:23,106 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:24,044 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:24,956 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:25,343 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:25,436 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:26,380 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:27,816 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:27,898 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:28,840 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:29,765 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:30,139 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:30,235 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:31,133 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:32,037 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:32,960 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:34,246 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:34,336 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:34,721 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:35,677 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:35,758 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:36,158 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:37,095 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:37,465 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:37,555 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:38,564 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:39,491 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:39,574 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:39,656 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:40,026 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:40,976 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:41,351 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:41,751 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:42,659 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:43,549 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:43,957 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:44,883 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:45,814 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:47,197 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:47,772 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:49,219 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:50,163 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:51,106 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:51,189 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:52,046 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:52,428 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:52,700 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:53,611 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:53,695 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:53,777 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:54,151 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:55,062 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:55,435 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:55,528 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:56,639 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:56,729 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:33:58,143 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:59,083 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:33:59,992 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:00,940 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:01,493 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:02,415 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:02,498 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:03,417 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:03,498 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:03,584 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:03,667 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:04,516 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:04,611 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:05,549 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:06,411 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:07,366 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:07,915 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:08,001 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:08,373 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:08,943 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:09,510 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:09,898 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:10,850 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:10,937 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:11,320 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:11,416 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:11,982 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:12,828 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:12,914 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:12,995 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:13,077 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:13,626 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:14,012 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:14,385 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:14,474 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:15,579 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:15,967 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:16,062 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:16,455 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:16,544 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:17,472 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:18,411 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:19,326 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:20,715 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:21,099 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:22,008 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:22,091 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:22,173 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:22,256 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:23,196 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:23,571 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:24,523 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:24,923 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:26,524 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:26,606 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:27,463 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:28,381 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:28,461 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:28,542 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:29,481 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:29,863 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:30,234 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:31,174 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:31,257 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:31,339 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:31,423 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:32,357 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:32,737 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:33,134 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:33,228 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:34,178 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:35,126 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:36,048 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:36,129 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:36,220 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:36,609 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:36,703 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:37,104 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:37,961 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:38,346 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:38,721 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:39,105 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:39,198 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:39,570 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:40,483 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:40,566 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:40,941 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:41,855 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:41,939 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:42,868 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:43,796 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:44,726 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:44,810 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:44,901 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:45,835 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:46,395 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:46,488 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:47,604 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:47,983 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:48,899 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:49,292 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:49,393 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:50,293 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:51,200 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:52,146 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:53,572 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:53,965 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:54,948 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:55,030 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:55,423 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:55,516 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:56,423 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:57,344 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:57,427 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:34:58,342 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:59,261 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:34:59,816 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:00,745 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:01,673 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:02,609 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:03,517 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:03,902 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:03,995 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:04,385 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:05,326 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:05,589 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:05,965 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:06,925 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:07,864 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:08,754 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:09,330 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:10,257 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:11,203 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:12,327 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:13,270 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:13,349 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:13,611 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:14,172 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:14,558 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:14,831 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:14,918 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:15,838 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:16,757 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:17,677 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:17,768 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:18,138 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:19,066 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:19,152 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:20,077 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:21,007 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:21,951 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:22,875 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:23,794 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:24,699 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:25,736 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:26,121 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:26,682 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:26,942 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:27,024 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:27,107 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:28,030 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:28,111 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:28,195 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:29,139 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:29,229 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:29,594 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:29,964 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:30,926 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:31,871 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:32,803 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:33,724 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:33,993 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:34,901 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:35,286 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:35,680 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:35,770 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:36,689 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:37,062 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:37,159 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:37,533 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:37,636 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:38,564 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:38,935 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:39,888 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:40,741 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:42,311 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:42,711 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:43,110 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:44,971 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:45,900 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:46,799 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:47,056 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:47,768 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:48,666 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:48,746 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:49,128 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:49,672 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:50,049 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:50,984 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:51,917 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:52,288 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:53,240 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:54,143 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:54,227 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:54,607 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:55,561 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:56,473 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:57,422 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:35:57,792 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:57,885 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:58,256 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:35:59,212 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:00,177 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:01,108 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:01,487 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:01,768 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:02,711 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:03,089 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:04,037 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:04,119 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:04,481 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:05,423 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:06,352 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:07,266 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:07,350 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:08,455 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:09,379 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:09,933 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:10,014 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:10,388 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:10,763 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:11,147 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:11,528 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:12,477 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:13,421 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:13,790 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:14,922 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:15,858 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:16,132 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:17,067 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:18,195 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:18,572 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:20,419 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:21,367 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:21,739 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:22,694 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:23,247 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:23,331 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:24,252 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:25,182 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:25,550 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:26,311 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:26,679 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:26,773 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:27,701 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:28,657 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:29,022 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:29,403 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:30,314 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:30,682 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:31,626 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:32,534 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:33,464 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:34,404 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:34,777 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:35,174 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:35,267 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:35,354 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:35,446 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:36,349 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:37,269 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:38,099 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:38,660 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:38,754 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:39,684 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:40,638 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:41,182 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:41,266 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:41,350 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:42,546 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:42,643 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:42,730 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:43,652 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:43,736 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:44,108 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:44,489 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:45,397 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:45,663 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:46,605 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:46,972 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:47,920 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:48,293 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:48,693 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:49,658 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:50,585 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:51,422 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:51,792 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:52,730 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:53,669 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:54,037 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:54,136 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:55,048 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:55,951 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:56,034 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:57,418 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:58,145 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:36:59,066 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:36:59,999 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:00,370 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:00,733 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:02,154 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:02,240 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:03,165 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:04,085 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:04,354 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:05,310 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:05,571 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:06,479 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:07,416 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:08,840 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:09,207 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:10,170 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:10,245 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:10,624 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:12,048 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:12,983 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:13,376 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:13,972 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:14,056 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:14,142 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:15,098 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:15,463 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:15,555 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:16,489 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:16,572 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:17,487 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:37:17,866 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:17,956 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:18,041 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:18,875 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:18,965 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:37:18,972 - distributed.scheduler - ERROR - Couldn't gather keys: {'run-47bf01a681f8c853d0bd517c5297f361': 'memory'}\n",
+      "2024-12-19 23:37:19,202 - distributed.scheduler - ERROR - Couldn't gather keys: {'run-47bf01a681f8c853d0bd517c5297f361': 'queued'}\n",
+      "2024-12-19 23:38:41,318 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:38:42,001 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:41:29,505 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:42:43,331 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:42:44,281 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:42:44,444 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:42:44,537 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:42:44,903 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:42:44,989 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:42:45,364 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:42:45,659 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-19 23:42:47,023 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-19 23:42:47,034 - distributed.scheduler - ERROR - Couldn't gather keys: {'run-90b1b2a59bcc7116d98d871a9a235520': 'memory'}\n",
+      "2024-12-19 23:42:47,255 - distributed.scheduler - ERROR - Couldn't gather keys: {'run-90b1b2a59bcc7116d98d871a9a235520': 'queued'}\n",
+      "2024-12-19 23:45:19,534 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:00:43,589 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:00:44,797 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:03:41,113 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:03:42,897 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:05:37,824 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:05:38,816 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:25,179 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:27,033 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:28,216 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:29,772 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:30,872 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:32,311 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:33,899 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:35,101 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:36,081 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:36,164 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:36,244 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:36,627 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:38,195 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:39,314 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:39,403 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:39,488 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:39,868 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:40,235 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:41,242 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:41,394 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:41,574 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:42,325 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:42,693 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:43,698 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:43,796 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:44,744 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:45,110 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:45,481 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:46,030 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:46,718 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:47,432 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:47,707 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:48,702 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:49,069 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:49,166 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:50,181 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:50,530 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:51,518 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:51,596 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:51,673 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:52,675 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:52,934 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:53,937 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:57,004 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:57,934 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:13:58,850 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:13:59,826 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:00,953 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:01,036 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:01,114 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:02,109 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:02,638 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:03,629 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:03,994 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:04,543 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:04,821 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:05,826 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:05,908 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:05,990 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:06,943 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:07,316 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:08,307 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:08,666 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:09,038 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:09,588 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:09,952 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:10,316 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:12,997 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:14,470 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:15,249 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:16,463 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:16,822 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:16,932 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:33,822 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:34,197 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:35,390 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:36,990 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:37,960 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:38,917 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:39,819 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:40,892 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:41,804 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:42,712 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:42,803 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:42,886 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:43,246 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:43,625 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:43,888 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:44,253 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:47,987 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:48,961 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:49,340 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:50,446 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:50,804 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:51,154 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:51,672 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:52,059 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:52,331 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:52,710 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:54,112 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:55,534 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:55,927 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:56,704 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:57,792 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:58,153 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:58,702 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:14:59,622 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:14:59,720 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:00,854 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:00,939 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:01,198 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:02,096 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:02,381 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:02,468 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:03,378 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:04,163 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:05,122 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:06,066 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:07,005 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:08,733 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:09,119 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:09,210 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:10,758 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:11,136 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:13,633 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:14,763 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:16,187 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:16,610 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:17,052 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:17,622 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:18,566 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:20,151 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:21,592 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:22,587 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:22,676 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:23,124 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:24,302 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:25,505 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:26,766 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:29,431 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:30,374 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:30,727 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:31,084 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:31,176 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:31,525 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:32,637 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:33,573 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:33,651 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:33,982 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:34,919 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:35,435 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:36,546 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:36,627 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:36,976 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:37,356 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:38,294 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:38,648 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:38,738 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:38,817 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:39,748 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:40,534 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:41,891 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:42,987 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:43,912 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:43,999 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:44,081 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:44,430 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:44,522 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:44,601 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:45,525 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:46,463 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:47,380 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:47,463 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:48,378 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:48,459 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:48,543 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:49,469 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:50,371 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:50,727 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:51,682 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:52,045 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:52,137 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:53,065 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:53,434 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:54,356 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:54,696 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:55,067 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:55,158 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:55,244 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:56,169 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:57,107 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:58,051 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:15:58,133 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:58,688 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:58,780 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:58,858 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:15:59,760 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:00,303 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:01,209 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:02,295 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:02,836 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:03,191 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:03,293 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:03,625 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:04,553 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:05,667 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:06,203 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:06,561 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:07,494 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:07,864 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:08,232 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:08,325 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:08,403 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:08,488 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:09,195 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:09,636 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:10,555 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:10,641 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:10,728 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:11,669 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:11,750 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:11,838 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:12,200 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:13,137 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:14,257 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:14,618 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:14,707 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:15,075 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:15,171 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:15,531 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:16,486 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:17,394 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:17,754 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:18,698 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:19,054 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:19,145 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:19,228 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:19,758 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:20,691 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:21,644 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:23,032 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:23,967 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:24,905 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:25,700 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:26,826 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:27,646 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:28,597 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:29,543 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:29,628 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:29,718 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:31,279 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:32,094 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:32,465 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:32,554 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:33,493 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:33,584 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:34,530 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:35,448 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:36,395 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:36,754 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:37,124 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:38,077 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:38,626 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:38,874 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:38,956 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:39,316 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:39,408 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:40,341 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:40,425 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:41,358 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:41,442 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:41,531 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:42,464 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:43,435 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:44,239 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:44,610 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:44,698 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:44,786 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:44,866 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:45,796 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:46,722 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:46,804 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:47,740 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:48,095 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:48,464 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:49,393 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:49,747 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:50,656 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:50,741 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:50,826 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:51,756 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:52,134 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:52,523 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:53,469 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:54,387 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:54,470 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:54,828 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:54,928 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:55,008 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:55,361 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:55,452 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:56,388 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:56,908 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:56,989 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:57,367 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:57,742 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:58,119 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:58,208 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:58,564 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:58,657 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:58,737 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:16:59,673 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:16:59,756 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:00,692 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:00,777 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:02,448 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:02,793 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:02,883 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:03,819 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:03,902 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:04,825 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:05,761 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:06,690 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:07,058 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:07,150 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:08,090 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:09,007 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:09,361 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:10,302 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:11,245 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:11,327 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:11,411 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:11,493 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:11,574 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:12,487 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:13,027 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:13,951 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:14,036 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:14,116 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:14,198 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:14,287 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:15,109 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:15,208 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:15,293 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:15,372 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:16,300 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:16,652 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:16,741 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:16,820 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:18,198 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:19,117 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:19,496 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:19,872 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:20,796 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:20,879 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:21,241 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:21,621 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:21,983 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:22,073 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:23,005 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:23,940 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:24,026 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:24,956 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:25,045 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:25,321 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:25,696 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:26,646 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:26,921 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:27,308 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:28,363 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:29,263 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:29,355 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:30,280 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:30,651 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:31,051 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:31,439 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:32,351 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:32,746 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:34,289 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:35,888 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:35,983 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:36,364 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:36,759 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:37,680 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "2024-12-20 00:17:38,076 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 111, in pickle_loads\n",
+      "    return pickle.loads(pik, buffers=buffers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 92, in loads\n",
+      "    return pickle.loads(x)\n",
+      "           ^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/storage.py\", line 371, in _load_from_bytes\n",
+      "    return torch.load(io.BytesIO(b))\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1040, in load\n",
+      "    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1268, in _legacy_load\n",
+      "    result = unpickler.load()\n",
+      "             ^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 1205, in persistent_load\n",
+      "    wrap_storage=restore_location(obj, location),\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 391, in default_restore_location\n",
+      "    result = fn(storage, location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 266, in _cuda_deserialize\n",
+      "    device = validate_cuda_device(location)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/torch/serialization.py\", line 257, in validate_cuda_device\n",
+      "    raise RuntimeError('Attempting to deserialize object on CUDA device '\n",
+      "RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.\n",
+      "2024-12-20 00:17:39,230 - distributed.protocol.core - CRITICAL - Failed to deserialize\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+      "    return msgpack.loads(\n",
+      "           ^^^^^^^^^^^^^^\n",
+      "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+      "    return merge_and_deserialize(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+      "    return func(*args, **kwds)\n",
+      "           ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+      "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+      "    return loads(header, frames)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+      "    raise TypeError(msg)\n",
+      "TypeError: Could not serialize object of type State\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+      "    result = pickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+      "    pickler.dump(x)\n",
+      "RecursionError: maximum recursion depth exceeded\n",
+      "\n",
+      "During handling of the above exception, another exception occurred:\n",
+      "\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+      "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+      "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+      "    frames[0] = pickle.dumps(\n",
+      "                ^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+      "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+      "    cp.dump(obj)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+      "    return super().dump(obj)\n",
+      "           ^^^^^^^^^^^^^^^^^\n",
+      "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "
00:17:39.232 | ERROR   | Flow run 'benign-parrot' - Encountered exception during execution: TypeError('Could not serialize object of type State\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\\n    result = pickle.dumps(x, **dump_kwargs)\\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\nTypeError: cannot pickle \\'weakref.ReferenceType\\' object\\n\\nDuring handling of the above exception, another exception occurred:\\n\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\\n    pickler.dump(x)\\nRecursionError: maximum recursion depth exceeded\\n\\nDuring handling of the above exception, another exception occurred:\\n\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\\n    header, frames = dumps(x, context=context) if wants_context else dumps(x)\\n                     ^^^^^^^^^^^^^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\\n    frames[0] = pickle.dumps(\\n                ^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\\n    result = cloudpickle.dumps(x, **dump_kwargs)\\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\\n    cp.dump(obj)\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\\n    return super().dump(obj)\\n           ^^^^^^^^^^^^^^^^^\\nTypeError: cannot pickle \\'weakref.ReferenceType\\' object\\n')\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 652, in run_context\n",
+       "    yield self\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 696, in run_flow_sync\n",
+       "    engine.call_flow_fn()\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 675, in call_flow_fn\n",
+       "    result = call_with_parameters(self.flow.fn, self.parameters)\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 206, in call_with_parameters\n",
+       "    return fn(*args, **kwargs)\n",
+       "           ^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/tasks/eos_alloy/flow.py\", line 130, in run_from_db\n",
+       "    return [\n",
+       "           ^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/tasks/eos_alloy/flow.py\", line 131, in <listcomp>\n",
+       "    f.result(timeout=None, raise_on_failure=False)\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py\", line 132, in result\n",
+       "    future_result = self._wrapped_future.result(timeout=timeout)\n",
+       "                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 402, in result\n",
+       "    return self.client.sync(self._result, callback_timeout=timeout)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/utils_comm.py\", line 441, in retry_operation\n",
+       "    return await retry(\n",
+       "           ^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/utils_comm.py\", line 420, in retry\n",
+       "    return await coro()\n",
+       "           ^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/core.py\", line 1259, in send_recv_from_rpc\n",
+       "    return await send_recv(comm=comm, op=key, **kwargs)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/core.py\", line 1018, in send_recv\n",
+       "    response = await comm.read(deserializers=deserializers)\n",
+       "               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/comm/tcp.py\", line 247, in read\n",
+       "    msg = await from_frames(\n",
+       "          ^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/comm/utils.py\", line 78, in from_frames\n",
+       "    res = _from_frames()\n",
+       "          ^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/comm/utils.py\", line 61, in _from_frames\n",
+       "    return protocol.loads(\n",
+       "           ^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+       "    return msgpack.loads(\n",
+       "           ^^^^^^^^^^^^^^\n",
+       "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+       "    return merge_and_deserialize(\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+       "    return func(*args, **kwds)\n",
+       "           ^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+       "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+       "    return loads(header, frames)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+       "    raise TypeError(msg)\n",
+       "TypeError: Could not serialize object of type State\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+       "    result = pickle.dumps(x, **dump_kwargs)\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+       "\n",
+       "During handling of the above exception, another exception occurred:\n",
+       "\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+       "    pickler.dump(x)\n",
+       "RecursionError: maximum recursion depth exceeded\n",
+       "\n",
+       "During handling of the above exception, another exception occurred:\n",
+       "\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+       "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+       "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+       "    frames[0] = pickle.dumps(\n",
+       "                ^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+       "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+       "    cp.dump(obj)\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+       "    return super().dump(obj)\n",
+       "           ^^^^^^^^^^^^^^^^^\n",
+       "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+       "\n",
+       "
\n"
+      ],
+      "text/plain": [
+       "00:17:39.232 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'benign-parrot'\u001b[0m - Encountered exception during execution: TypeError('Could not serialize object of type State\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\\n    result = pickle.dumps(x, **dump_kwargs)\\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\nTypeError: cannot pickle \\'weakref.ReferenceType\\' object\\n\\nDuring handling of the above exception, another exception occurred:\\n\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\\n    pickler.dump(x)\\nRecursionError: maximum recursion depth exceeded\\n\\nDuring handling of the above exception, another exception occurred:\\n\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\\n    header, frames = dumps(x, context=context) if wants_context else dumps(x)\\n                     ^^^^^^^^^^^^^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\\n    frames[0] = pickle.dumps(\\n                ^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\\n    result = cloudpickle.dumps(x, **dump_kwargs)\\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\\n    cp.dump(obj)\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\\n    return super().dump(obj)\\n           ^^^^^^^^^^^^^^^^^\\nTypeError: cannot pickle \\'weakref.ReferenceType\\' object\\n')\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 652, in run_context\n",
+       "    yield self\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 696, in run_flow_sync\n",
+       "    engine.call_flow_fn()\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 675, in call_flow_fn\n",
+       "    result = call_with_parameters(self.flow.fn, self.parameters)\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 206, in call_with_parameters\n",
+       "    return fn(*args, **kwargs)\n",
+       "           ^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/tasks/eos_alloy/flow.py\", line 130, in run_from_db\n",
+       "    return [\n",
+       "           ^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/tasks/eos_alloy/flow.py\", line 131, in \n",
+       "    f.result(timeout=None, raise_on_failure=False)\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py\", line 132, in result\n",
+       "    future_result = self._wrapped_future.result(timeout=timeout)\n",
+       "                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 402, in result\n",
+       "    return self.client.sync(self._result, callback_timeout=timeout)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/utils_comm.py\", line 441, in retry_operation\n",
+       "    return await retry(\n",
+       "           ^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/utils_comm.py\", line 420, in retry\n",
+       "    return await coro()\n",
+       "           ^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/core.py\", line 1259, in send_recv_from_rpc\n",
+       "    return await send_recv(comm=comm, op=key, **kwargs)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/core.py\", line 1018, in send_recv\n",
+       "    response = await comm.read(deserializers=deserializers)\n",
+       "               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/comm/tcp.py\", line 247, in read\n",
+       "    msg = await from_frames(\n",
+       "          ^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/comm/utils.py\", line 78, in from_frames\n",
+       "    res = _from_frames()\n",
+       "          ^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/comm/utils.py\", line 61, in _from_frames\n",
+       "    return protocol.loads(\n",
+       "           ^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 175, in loads\n",
+       "    return msgpack.loads(\n",
+       "           ^^^^^^^^^^^^^^\n",
+       "  File \"msgpack/_unpacker.pyx\", line 194, in msgpack._cmsgpack.unpackb\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py\", line 159, in _decode_default\n",
+       "    return merge_and_deserialize(\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py\", line 81, in inner\n",
+       "    return func(*args, **kwds)\n",
+       "           ^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 525, in merge_and_deserialize\n",
+       "    return deserialize(header, merged_frames, deserializers=deserializers)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 452, in deserialize\n",
+       "    return loads(header, frames)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 195, in serialization_error_loads\n",
+       "    raise TypeError(msg)\n",
+       "TypeError: Could not serialize object of type State\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n",
+       "    result = pickle.dumps(x, **dump_kwargs)\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+       "\n",
+       "During handling of the above exception, another exception occurred:\n",
+       "\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n",
+       "    pickler.dump(x)\n",
+       "RecursionError: maximum recursion depth exceeded\n",
+       "\n",
+       "During handling of the above exception, another exception occurred:\n",
+       "\n",
+       "Traceback (most recent call last):\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n",
+       "    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n",
+       "                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n",
+       "    frames[0] = pickle.dumps(\n",
+       "                ^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n",
+       "    result = cloudpickle.dumps(x, **dump_kwargs)\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n",
+       "    cp.dump(obj)\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n",
+       "    return super().dump(obj)\n",
+       "           ^^^^^^^^^^^^^^^^^\n",
+       "TypeError: cannot pickle 'weakref.ReferenceType' object\n",
+       "\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "
00:17:39.489 | ERROR   | Flow run 'benign-parrot' - Finished in state Failed('Flow run encountered an exception: TypeError: Could not serialize object of type State\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\\n    result = pickle.dumps(x, **dump_kwargs)\\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\nTypeError: cannot pickle \\'weakref.ReferenceType\\' object\\n\\nDuring handling of the above exception, another exception occurred:\\n\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\\n    pickler.dump(x)\\nRecursionError: maximum recursion depth exceeded\\n\\nDuring handling of the above exception, another exception occurred:\\n\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\\n    header, frames = dumps(x, context=context) if wants_context else dumps(x)\\n                     ^^^^^^^^^^^^^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\\n    frames[0] = pickle.dumps(\\n                ^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\\n    result = cloudpickle.dumps(x, **dump_kwargs)\\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\\n    cp.dump(obj)\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\\n    return super().dump(obj)\\n           ^^^^^^^^^^^^^^^^^\\nTypeError: cannot pickle \\'weakref.ReferenceType\\' object\\n')\n",
+       "
\n"
+      ],
+      "text/plain": [
+       "00:17:39.489 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'benign-parrot'\u001b[0m - Finished in state \u001b[38;5;160mFailed\u001b[0m('Flow run encountered an exception: TypeError: Could not serialize object of type State\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\\n    result = pickle.dumps(x, **dump_kwargs)\\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\nTypeError: cannot pickle \\'weakref.ReferenceType\\' object\\n\\nDuring handling of the above exception, another exception occurred:\\n\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\\n    pickler.dump(x)\\nRecursionError: maximum recursion depth exceeded\\n\\nDuring handling of the above exception, another exception occurred:\\n\\nTraceback (most recent call last):\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\\n    header, frames = dumps(x, context=context) if wants_context else dumps(x)\\n                     ^^^^^^^^^^^^^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\\n    frames[0] = pickle.dumps(\\n                ^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\\n    result = cloudpickle.dumps(x, **dump_kwargs)\\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\\n    cp.dump(obj)\\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\\n    return super().dump(obj)\\n           ^^^^^^^^^^^^^^^^^\\nTypeError: cannot pickle \\'weakref.ReferenceType\\' object\\n')\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "ename": "TypeError",
+     "evalue": "Could not serialize object of type State\nTraceback (most recent call last):\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n    result = pickle.dumps(x, **dump_kwargs)\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nTypeError: cannot pickle 'weakref.ReferenceType' object\n\nDuring handling of the above exception, another exception occurred:\n\nTraceback (most recent call last):\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n    pickler.dump(x)\nRecursionError: maximum recursion depth exceeded\n\nDuring handling of the above exception, another exception occurred:\n\nTraceback (most recent call last):\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n    frames[0] = pickle.dumps(\n                ^^^^^^^^^^^^^\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n    result = cloudpickle.dumps(x, **dump_kwargs)\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n    cp.dump(obj)\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n    return super().dump(obj)\n           ^^^^^^^^^^^^^^^^^\nTypeError: cannot pickle 'weakref.ReferenceType' object\n",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[2], line 41\u001b[0m\n\u001b[1;32m     34\u001b[0m run_from_db_ \u001b[38;5;241m=\u001b[39m run_from_db\u001b[38;5;241m.\u001b[39mwith_options(\n\u001b[1;32m     35\u001b[0m     task_runner\u001b[38;5;241m=\u001b[39mDaskTaskRunner(address\u001b[38;5;241m=\u001b[39mclient\u001b[38;5;241m.\u001b[39mscheduler\u001b[38;5;241m.\u001b[39maddress),\n\u001b[1;32m     36\u001b[0m     \u001b[38;5;66;03m# log_prints=True,\u001b[39;00m\n\u001b[1;32m     37\u001b[0m )\n\u001b[1;32m     39\u001b[0m \u001b[38;5;66;03m# Run the workflow\u001b[39;00m\n\u001b[0;32m---> 41\u001b[0m results \u001b[38;5;241m=\u001b[39m \u001b[43mrun_from_db_\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     42\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdb_path\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43msqs_Fe-Ni-Cr.db\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_path\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43meos.h5\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtable_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mFe-Ni-Cr\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\n\u001b[1;32m     43\u001b[0m \u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flows.py:1345\u001b[0m, in \u001b[0;36mFlow.__call__\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1341\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m track_viz_task(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39misasync, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mname, parameters)\n\u001b[1;32m   1343\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mprefect\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mflow_engine\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m run_flow\n\u001b[0;32m-> 1345\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_flow\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1346\u001b[0m \u001b[43m    \u001b[49m\u001b[43mflow\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1347\u001b[0m \u001b[43m    \u001b[49m\u001b[43mparameters\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1348\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwait_for\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwait_for\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1349\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_type\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1350\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:818\u001b[0m, in \u001b[0;36mrun_flow\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    816\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m run_flow_async(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    817\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 818\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_flow_sync\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:698\u001b[0m, in \u001b[0;36mrun_flow_sync\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    695\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mrun_context():\n\u001b[1;32m    696\u001b[0m             engine\u001b[38;5;241m.\u001b[39mcall_flow_fn()\n\u001b[0;32m--> 698\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;28;01mif\u001b[39;00m return_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstate\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:255\u001b[0m, in \u001b[0;36mFlowRunEngine.result\u001b[0;34m(self, raise_on_failure)\u001b[0m\n\u001b[1;32m    253\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m NotSet:\n\u001b[1;32m    254\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m raise_on_failure:\n\u001b[0;32m--> 255\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised\n\u001b[1;32m    256\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised\n\u001b[1;32m    258\u001b[0m \u001b[38;5;66;03m# This is a fall through case which leans on the existing state result mechanics to get the\u001b[39;00m\n\u001b[1;32m    259\u001b[0m \u001b[38;5;66;03m# return value. This is necessary because we currently will return a State object if the\u001b[39;00m\n\u001b[1;32m    260\u001b[0m \u001b[38;5;66;03m# the State was Prefect-created.\u001b[39;00m\n\u001b[1;32m    261\u001b[0m \u001b[38;5;66;03m# TODO: Remove the need to get the result from a State except in cases where the return value\u001b[39;00m\n\u001b[1;32m    262\u001b[0m \u001b[38;5;66;03m# is a State object.\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:652\u001b[0m, in \u001b[0;36mFlowRunEngine.run_context\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    645\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m timeout_context(\n\u001b[1;32m    646\u001b[0m         seconds\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow\u001b[38;5;241m.\u001b[39mtimeout_seconds,\n\u001b[1;32m    647\u001b[0m         timeout_exc_type\u001b[38;5;241m=\u001b[39mFlowRunTimeoutError,\n\u001b[1;32m    648\u001b[0m     ):\n\u001b[1;32m    649\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\n\u001b[1;32m    650\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mExecuting flow \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m for flow run \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow_run\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m...\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    651\u001b[0m         )\n\u001b[0;32m--> 652\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n\u001b[1;32m    653\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    654\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandle_timeout(exc)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:696\u001b[0m, in \u001b[0;36mrun_flow_sync\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    694\u001b[0m     \u001b[38;5;28;01mwhile\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mis_running():\n\u001b[1;32m    695\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mrun_context():\n\u001b[0;32m--> 696\u001b[0m             \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_flow_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    698\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;28;01mif\u001b[39;00m return_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstate\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mresult()\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:675\u001b[0m, in \u001b[0;36mFlowRunEngine.call_flow_fn\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    673\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _call_flow_fn()\n\u001b[1;32m    674\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 675\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[43mcall_with_parameters\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mflow\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    676\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandle_success(result)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py:206\u001b[0m, in \u001b[0;36mcall_with_parameters\u001b[0;34m(fn, parameters)\u001b[0m\n\u001b[1;32m    198\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;124;03mCall a function with parameters extracted with `get_call_parameters`\u001b[39;00m\n\u001b[1;32m    200\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    203\u001b[0m \u001b[38;5;124;03mthe args/kwargs using `parameters_to_positional_and_keyword` directly\u001b[39;00m\n\u001b[1;32m    204\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    205\u001b[0m args, kwargs \u001b[38;5;241m=\u001b[39m parameters_to_args_kwargs(fn, parameters)\n\u001b[0;32m--> 206\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/tasks/eos_alloy/flow.py:130\u001b[0m, in \u001b[0;36mrun_from_db\u001b[0;34m(db_path, out_path, table_name, optimizer, optimizer_kwargs, filter, filter_kwargs, criterion, max_abs_strain, concurrent)\u001b[0m\n\u001b[1;32m    126\u001b[0m         futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[1;32m    128\u001b[0m wait(futures)\n\u001b[0;32m--> 130\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m[\u001b[49m\n\u001b[1;32m    131\u001b[0m \u001b[43m    \u001b[49m\u001b[43mf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mraise_on_failure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m    132\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mf\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfutures\u001b[49m\n\u001b[1;32m    133\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstate\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mis_completed\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    134\u001b[0m \u001b[43m\u001b[49m\u001b[43m]\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/mlip-arena/mlip_arena/tasks/eos_alloy/flow.py:131\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    126\u001b[0m         futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[1;32m    128\u001b[0m wait(futures)\n\u001b[1;32m    130\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m [\n\u001b[0;32m--> 131\u001b[0m     \u001b[43mf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mraise_on_failure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m    132\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m f \u001b[38;5;129;01min\u001b[39;00m futures\n\u001b[1;32m    133\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m f\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mis_completed()\n\u001b[1;32m    134\u001b[0m ]\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py:132\u001b[0m, in \u001b[0;36mPrefectDaskFuture.result\u001b[0;34m(self, timeout, raise_on_failure)\u001b[0m\n\u001b[1;32m    130\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_final_state:\n\u001b[1;32m    131\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 132\u001b[0m         future_result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_wrapped_future\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    133\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m distributed\u001b[38;5;241m.\u001b[39mTimeoutError \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    134\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m(\n\u001b[1;32m    135\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTask run \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtask_run_id\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m did not complete within \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtimeout\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m seconds\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    136\u001b[0m         ) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mexc\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py:402\u001b[0m, in \u001b[0;36mFuture.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    400\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_verify_initialized()\n\u001b[1;32m    401\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m shorten_traceback():\n\u001b[0;32m--> 402\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mclient\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msync\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_result\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallback_timeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/utils_comm.py:441\u001b[0m, in \u001b[0;36mretry_operation\u001b[0;34m(coro, operation, *args, **kwargs)\u001b[0m\n\u001b[1;32m    435\u001b[0m retry_delay_min \u001b[38;5;241m=\u001b[39m parse_timedelta(\n\u001b[1;32m    436\u001b[0m     dask\u001b[38;5;241m.\u001b[39mconfig\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdistributed.comm.retry.delay.min\u001b[39m\u001b[38;5;124m\"\u001b[39m), default\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ms\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    437\u001b[0m )\n\u001b[1;32m    438\u001b[0m retry_delay_max \u001b[38;5;241m=\u001b[39m parse_timedelta(\n\u001b[1;32m    439\u001b[0m     dask\u001b[38;5;241m.\u001b[39mconfig\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdistributed.comm.retry.delay.max\u001b[39m\u001b[38;5;124m\"\u001b[39m), default\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ms\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    440\u001b[0m )\n\u001b[0;32m--> 441\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m retry(\n\u001b[1;32m    442\u001b[0m     partial(coro, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs),\n\u001b[1;32m    443\u001b[0m     count\u001b[38;5;241m=\u001b[39mretry_count,\n\u001b[1;32m    444\u001b[0m     delay_min\u001b[38;5;241m=\u001b[39mretry_delay_min,\n\u001b[1;32m    445\u001b[0m     delay_max\u001b[38;5;241m=\u001b[39mretry_delay_max,\n\u001b[1;32m    446\u001b[0m     operation\u001b[38;5;241m=\u001b[39moperation,\n\u001b[1;32m    447\u001b[0m )\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/utils_comm.py:420\u001b[0m, in \u001b[0;36mretry\u001b[0;34m(coro, count, delay_min, delay_max, jitter_fraction, retry_on_exceptions, operation)\u001b[0m\n\u001b[1;32m    418\u001b[0m             delay \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m \u001b[38;5;241m+\u001b[39m random\u001b[38;5;241m.\u001b[39mrandom() \u001b[38;5;241m*\u001b[39m jitter_fraction\n\u001b[1;32m    419\u001b[0m         \u001b[38;5;28;01mawait\u001b[39;00m asyncio\u001b[38;5;241m.\u001b[39msleep(delay)\n\u001b[0;32m--> 420\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m coro()\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/core.py:1259\u001b[0m, in \u001b[0;36mPooledRPCCall.__getattr__..send_recv_from_rpc\u001b[0;34m(**kwargs)\u001b[0m\n\u001b[1;32m   1257\u001b[0m prev_name, comm\u001b[38;5;241m.\u001b[39mname \u001b[38;5;241m=\u001b[39m comm\u001b[38;5;241m.\u001b[39mname, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mConnectionPool.\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m key\n\u001b[1;32m   1258\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1259\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m send_recv(comm\u001b[38;5;241m=\u001b[39mcomm, op\u001b[38;5;241m=\u001b[39mkey, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m   1260\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m   1261\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpool\u001b[38;5;241m.\u001b[39mreuse(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maddr, comm)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/core.py:1018\u001b[0m, in \u001b[0;36msend_recv\u001b[0;34m(comm, reply, serializers, deserializers, **kwargs)\u001b[0m\n\u001b[1;32m   1016\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m comm\u001b[38;5;241m.\u001b[39mwrite(msg, serializers\u001b[38;5;241m=\u001b[39mserializers, on_error\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mraise\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m   1017\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m reply:\n\u001b[0;32m-> 1018\u001b[0m     response \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m comm\u001b[38;5;241m.\u001b[39mread(deserializers\u001b[38;5;241m=\u001b[39mdeserializers)\n\u001b[1;32m   1019\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1020\u001b[0m     response \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/comm/tcp.py:247\u001b[0m, in \u001b[0;36mTCP.read\u001b[0;34m(self, deserializers)\u001b[0m\n\u001b[1;32m    245\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    246\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 247\u001b[0m         msg \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m from_frames(\n\u001b[1;32m    248\u001b[0m             frames,\n\u001b[1;32m    249\u001b[0m             deserialize\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdeserialize,\n\u001b[1;32m    250\u001b[0m             deserializers\u001b[38;5;241m=\u001b[39mdeserializers,\n\u001b[1;32m    251\u001b[0m             allow_offload\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mallow_offload,\n\u001b[1;32m    252\u001b[0m         )\n\u001b[1;32m    253\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mEOFError\u001b[39;00m:\n\u001b[1;32m    254\u001b[0m         \u001b[38;5;66;03m# Frames possibly garbled or truncated by communication error\u001b[39;00m\n\u001b[1;32m    255\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mabort()\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/comm/utils.py:78\u001b[0m, in \u001b[0;36mfrom_frames\u001b[0;34m(frames, deserialize, deserializers, allow_offload)\u001b[0m\n\u001b[1;32m     76\u001b[0m     res \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m offload(_from_frames)\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m---> 78\u001b[0m     res \u001b[38;5;241m=\u001b[39m \u001b[43m_from_frames\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m res\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/comm/utils.py:61\u001b[0m, in \u001b[0;36mfrom_frames.._from_frames\u001b[0;34m()\u001b[0m\n\u001b[1;32m     59\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_from_frames\u001b[39m():\n\u001b[1;32m     60\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 61\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mprotocol\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mloads\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     62\u001b[0m \u001b[43m            \u001b[49m\u001b[43mframes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdeserialize\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdeserialize\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdeserializers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdeserializers\u001b[49m\n\u001b[1;32m     63\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     64\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mEOFError\u001b[39;00m:\n\u001b[1;32m     65\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m size \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m1000\u001b[39m:\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py:175\u001b[0m, in \u001b[0;36mloads\u001b[0;34m(frames, deserialize, deserializers)\u001b[0m\n\u001b[1;32m    172\u001b[0m             \u001b[38;5;28;01mreturn\u001b[39;00m pickle\u001b[38;5;241m.\u001b[39mloads(sub_header[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpickled-obj\u001b[39m\u001b[38;5;124m\"\u001b[39m], buffers\u001b[38;5;241m=\u001b[39msub_frames)\n\u001b[1;32m    173\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m msgpack_decode_default(obj)\n\u001b[0;32m--> 175\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmsgpack\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mloads\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    176\u001b[0m \u001b[43m        \u001b[49m\u001b[43mframes\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mobject_hook\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m_decode_default\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43muse_list\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mmsgpack_opts\u001b[49m\n\u001b[1;32m    177\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    179\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n\u001b[1;32m    180\u001b[0m     logger\u001b[38;5;241m.\u001b[39mcritical(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFailed to deserialize\u001b[39m\u001b[38;5;124m\"\u001b[39m, exc_info\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
+      "File \u001b[0;32mmsgpack/_unpacker.pyx:194\u001b[0m, in \u001b[0;36mmsgpack._cmsgpack.unpackb\u001b[0;34m()\u001b[0m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/core.py:159\u001b[0m, in \u001b[0;36mloads.._decode_default\u001b[0;34m(obj)\u001b[0m\n\u001b[1;32m    157\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcompression\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m sub_header:\n\u001b[1;32m    158\u001b[0m         sub_frames \u001b[38;5;241m=\u001b[39m decompress(sub_header, sub_frames)\n\u001b[0;32m--> 159\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmerge_and_deserialize\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    160\u001b[0m \u001b[43m        \u001b[49m\u001b[43msub_header\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msub_frames\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdeserializers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdeserializers\u001b[49m\n\u001b[1;32m    161\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    162\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    163\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m Serialized(sub_header, sub_frames)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/contextlib.py:81\u001b[0m, in \u001b[0;36mContextDecorator.__call__..inner\u001b[0;34m(*args, **kwds)\u001b[0m\n\u001b[1;32m     78\u001b[0m \u001b[38;5;129m@wraps\u001b[39m(func)\n\u001b[1;32m     79\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minner\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwds):\n\u001b[1;32m     80\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_recreate_cm():\n\u001b[0;32m---> 81\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwds\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py:525\u001b[0m, in \u001b[0;36mmerge_and_deserialize\u001b[0;34m(header, frames, deserializers)\u001b[0m\n\u001b[1;32m    521\u001b[0m             merged \u001b[38;5;241m=\u001b[39m host_array_from_buffers(subframes)\n\u001b[1;32m    523\u001b[0m         merged_frames\u001b[38;5;241m.\u001b[39mappend(merged)\n\u001b[0;32m--> 525\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mdeserialize\u001b[49m\u001b[43m(\u001b[49m\u001b[43mheader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmerged_frames\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdeserializers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdeserializers\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py:452\u001b[0m, in \u001b[0;36mdeserialize\u001b[0;34m(header, frames, deserializers)\u001b[0m\n\u001b[1;32m    447\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\n\u001b[1;32m    448\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mData serialized with \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m but only able to deserialize \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    449\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdata with \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m%\u001b[39m (name, \u001b[38;5;28mstr\u001b[39m(\u001b[38;5;28mlist\u001b[39m(deserializers)))\n\u001b[1;32m    450\u001b[0m     )\n\u001b[1;32m    451\u001b[0m dumps, loads, wants_context \u001b[38;5;241m=\u001b[39m families[name]\n\u001b[0;32m--> 452\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mloads\u001b[49m\u001b[43m(\u001b[49m\u001b[43mheader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mframes\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py:195\u001b[0m, in \u001b[0;36mserialization_error_loads\u001b[0;34m(header, frames)\u001b[0m\n\u001b[1;32m    193\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mserialization_error_loads\u001b[39m(header, frames):\n\u001b[1;32m    194\u001b[0m     msg \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mjoin([codecs\u001b[38;5;241m.\u001b[39mdecode(frame, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mutf8\u001b[39m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;28;01mfor\u001b[39;00m frame \u001b[38;5;129;01min\u001b[39;00m frames])\n\u001b[0;32m--> 195\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(msg)\n",
+      "\u001b[0;31mTypeError\u001b[0m: Could not serialize object of type State\nTraceback (most recent call last):\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 60, in dumps\n    result = pickle.dumps(x, **dump_kwargs)\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nTypeError: cannot pickle 'weakref.ReferenceType' object\n\nDuring handling of the above exception, another exception occurred:\n\nTraceback (most recent call last):\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 65, in dumps\n    pickler.dump(x)\nRecursionError: maximum recursion depth exceeded\n\nDuring handling of the above exception, another exception occurred:\n\nTraceback (most recent call last):\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 366, in serialize\n    header, frames = dumps(x, context=context) if wants_context else dumps(x)\n                     ^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/serialize.py\", line 78, in pickle_dumps\n    frames[0] = pickle.dumps(\n                ^^^^^^^^^^^^^\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/protocol/pickle.py\", line 77, in dumps\n    result = cloudpickle.dumps(x, **dump_kwargs)\n             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1529, in dumps\n    cp.dump(obj)\n  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/cloudpickle/cloudpickle.py\", line 1295, in dump\n    return super().dump(obj)\n           ^^^^^^^^^^^^^^^^^\nTypeError: cannot pickle 'weakref.ReferenceType' object\n"
+     ]
     }
    ],
    "source": [
@@ -189,7 +54074,7 @@
     "    memory=\"64 GB\",\n",
     "    shebang=\"#!/bin/bash\",\n",
     "    account=\"m3828\",\n",
-    "    walltime=\"00:30:00\",\n",
+    "    walltime=\"01:00:00\",\n",
     "    job_mem=\"0\",\n",
     "    job_script_prologue=[\n",
     "        \"source ~/.bashrc\",\n",
@@ -208,7 +54093,7 @@
     "\n",
     "cluster = SLURMCluster(**cluster_kwargs)\n",
     "print(cluster.job_script())\n",
-    "cluster.adapt(minimum_jobs=2, maximum_jobs=2)\n",
+    "cluster.adapt(minimum_jobs=10, maximum_jobs=30)\n",
     "client = Client(cluster)\n",
     "\n",
     "run_from_db_ = run_from_db.with_options(\n",
@@ -12694,7 +66579,7 @@
         }
        }
       },
-      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAeQAAAFoCAYAAACG3IhaAAAAAXNSR0IArs4c6QAAIABJREFUeF7sXQd4k9X3frOb0kmhlKEIiIKKiAICiiAoCCJLQUFEhgwZAoKC7K0Iguw9BJUpyhQUEMUBLtSfCxUVZBa690jz/5+bJk2atBnN+L703OfxweSuc9/7Ne937j1DYTQajeDCCDACjAAjwAgwAgFFQMGEHFD8eXJGgBFgBBgBRkAgwITMDwIjwAgwAowAIyABBJiQJbAJLAIjwAgwAowAI8CEzM8AI8AIMAKMACMgAQSYkCWwCSwCI8AIMAKMACPAhMzPACPACDACjAAjIAEEmJAlsAksAiPACDACjAAjwITMzwAjwAgwAowAIyABBJiQJbAJLAIjwAgwAowAI8CEzM8AI8AIMAKMACMgAQSYkCWwCSwCI8AIMAKMACPAhMzPACPACDACjAAjIAEEmJAlsAksAiPACDACjAAjwITMzwAjwAgwAowAIyABBJiQJbAJLAIjwAgwAowAI8CE7MVngFJLJyanITUtA1GRYYgMD4NSqbCZId9gQHZ2LrQaNbRajRdn56EYAUaAEWAE5IyAXwj57LlL6PzsRIHT+OG90LdHezvMXp61CgeOnhTf/3R0A1QqpVu4bt9zDJfjEzF60BM2/QyGAry5dhdq16yKbh1aujWmq42zsnOxaceH2LD1Q2RmZdt0u6/JHejxWGs8/EBj8f2+j77EhLlrMOjpTnayujoft2MEGAFGgBEIPgT8Qsh//XMRXfpPEuhVjArHsZ2LoNGoLWheunIdDz81zvL5x6ProVap3EL7mZFz8f3//sAvxzfZ9MvLy8ddDz+H1i3uwvK5o90a09XGi9bsxLp3D4i1tW7RCLfUroHzF6/i9M9/4bc/z6HlvQ2wat5YMdzJ73/FWzsOo12rxj57QXBVbm7HCDACjAAjIB0E/ErIofoQoUG+NnEwHmvXwoLCwtU7sH7rwYATMh05KxS2R8zOtsr8skEk/O6KqdCHaG26HD7+jSDhaS8+62wop/WeyOd0UG7ACDACjAAjIAkE/ErI/Z/qgD2HPkd0ZDj2bJojyC89IwsPPjEGDz9wD/74+4LQKK015PkrtuHbH8/g0tXr4n62RtXKeLhVY/R/sgNioiMEiHMWb8EHh74QZP9As4YWYMcO7YHXl2/DF9/8DHoZaNzwVlGnD9Fh4fRh4v9T0jKwbMNu0ebchatCu6Wj7ae7P2w5Nv/lzL9YtvF9PNWlDW6oHov9H3+JP/++gOaN70BkRAXQcfszT7TDhBG9nW4qrW/J+t3o2bk1HmzRSLSft3wrrl5LwtC+nbFs42589e2vQtvu3vEBDHnmMZz++U+seXu/wCFEpxHfjxzY3e1TBKfCcQNGgBFgBBiBgCHgV0J+vm8XhIRoQUe8q+a9iJb33ol33z8qCHXX2hmY8voGO0Ju2XUksnPy0PD2OogIq4BfzvyDC5ev4Y5ba+GdFZMFKVnfP9evW9MC5qyXB2Diq2sF0VMx14VV0GPTmxOQkJSKHoOnCTIkom94Wx3LPfaApzpi7NCeot+JUz9h6PiFoj8RqrmQlj9q4ON46Mmxoj+NWbVKTKmb+dW3v+C5cfMx8YU+eLr7Q6Ltk0Nm4Ocz/1j60drMn+9tVB+nTv8m6qy/XzpnFNrcZyJ0LowAI8AIMALyR8DvhEyaZIvOw9HkrnpY98ZLaPfUONSuWQ3rFryEJwZNsyNkIsC6tWtYtMGCAiNGTVmCY1+cxp6Nc3BzrepiFzy5Q561aDO27TmGua8MQud2LYTGTlo2yUHa8qe7F6NSxUgLIdM8z/V+FI882BSxlaKRm5ePqrEVLXNTffvWTdGgXi3Uu/lG3FGvFsLDQm2ektIIeXj/bhjYqyN0Wg2s791HDOgGekGg78+c/Q/dB04RZE6kzoURYAQYAUYgOBDwOyETucxd8g7e2f0x6Ah747YPsWb+OJA1siNCJpiJhP85fwkXLtOxdSo++fI0jp74XhhpkbGWJ4RMYzZo019otgfefg0KFN0dr3jrA6zavBcbF01A00b1LIQ8buiTQubihTTs2W9uFi8JxQtZffd5vJ3lbrkkQv77/GV88+Eqm+5teowBGaWd+GCpzfd0akBEf/DtecHxFPIqGAFGgBFgBBAQQiYL5A5Pjxfw16lZzXKf7IiQj5z4DjPe2CTuj4sX62NbdzVkcpF6qOeLpT4C8yYNQaeHm1sIeeqYvniyS5sS+9AR+M+//4MzZ88LC+vPTv4o2j7athlenzJU/L87hNxtwGTxElKcqDv2GY+09Ew7oubnmRFgBBgBRkC+CASEkAmu0VOX4ePPvhXHxV3a3ycQLE7IZvIigyzSrO+sXxvV4yrj6OffYfabW1AWQjb7RtO9LPkJOyqkHd9YvYrLhFx8jLP/XkTnfiZ3r28PrRFasjcImXy6k1LSmJDl+3fHkjMCjAAjYIdAwAj5z38u4OPPvsNzvTpaIlYVJ2SyPt6887Dw4SVfXnN5/8MTmDxvvVuEbO0LTONk5+TinvaD0eiOunh7mYk0Sypmoy5HGjJF3irNZ7rf6NfwzQ+/Y++mOahzU3UmZP4jZAQYAUaAEXCIQMAI2ZE0xQl53MyV+PDYKaxf+DKa3X2b6EIE+PryrXhn9xEbQn5hyhJxr/zJrjcRWynKZvjbW/dDzRpV7O5cew2bhZ9+PWux+LbuRK5GpB2Ta1VphLxp+yFhEf3i4B6oFlfJZt7LVxOEhkyGYj8cWQ+NWsWEzH+IjAAjwAgwAvIjZAqHOXPRZmF49djDZAUN0J2y2Y3J+sh67Tv7RYhMst5u37qJcGWi+16ygiY3IzoqJv/d226piUtXEoRLE/kX9xwyXQBDPsZkFX0tIVn4+5JfMrlikauTM0Kev3KbGIN8oG+/5Sbh2vXdT39Y7pBnvjQAjz/6gGjDR9b8l8gIMAKMACPgCAG/aMjmu9Rhz3YBufaUVMwasjmWNWnD0+ZvxAeHPrd0ufO2OsKlaMfeT7Bs7ihLcA0KMLJ43S5LgBDqYD4mpvmXb/oAFDWLCt1Jmw2lSLt9bem7IviGdSFDrPEjegsNmch58EsLRLStnp0ftGl3/mI8du0/DjpGL254RgZrZGXd5v67LX1OfvcrBo59HZNGPYPe3dqK78kP2ZGVtdn9qrhRFxl7XU9M4Ttk/ptmBBgBRiCIEPALIZcVr/jryUJzjakYgbjKFUsdjtyELscnCP9hIl7rkpqeKTIxVakUbRNLm9rQnfKlqwnQ67SoXCnKoyhYZPl8LTEFBYYCcXwdqteVdencnxFgBBgBRqCcICALQi4ne8HLZAQYAUaAESjHCDAhl+PN56UzAowAI8AISAcBJmTp7AVLwggwAowAI1COEWBCLsebz0tnBBgBRoARkA4CTMjS2QuWhBGwQYB85MmavnhRqVRo1bwozSjDxggwAsGBABNycOwjryIIETAHuym+NGu3vSBcNi+JESi3CDAhl9ut54VLHQEi5KvxSXh7+WQbUSlATmnhWqW+LpaPEWAEHCPAhMxPBiMgUQSIkK8lpGDriikOJSS/9yXr38PRz78XkenubVQfLw/vJQLncGEEGAH5IcCELL89Y4nLCQJEyH/9cxGzxz9ns+J6N98AnVaL3sNmITk1Hb27P4SKkeF4+72PRcS3YzsXinzZXBgBRkBeCDAhy2u/WNpyhEBJd8g7Vk9HfEISRkxcjHdXTEHD2+oIVCjGO4VVXTxrJB5qeU85QoqXyggEBwJMyMGxj7yKIESACJkSoaya96LN6qIiw7DunQNYumG3SH5iLgaDQZDy+OG90LdH+yBEhJfECAQ3AkzIwb2/vDoZI1DaHTJlNqMMZ8XJmpZbs0YcbqweK+OVs+iMQPlEgAm5fO47r1oGCJRGyHsOf4GJr67Fno1zcHOt6jarMRqNUJApNhdGgBGQFQJMyLLaLha2PCFQGiFnZGbjsWdfQYhOi/HDe+OmG+Lw739XsOfw53isXQtLWtLyhBevlRGQOwJMyHLfQZY/aBFw5vZEFtWzF23GqdO/WTCgO+U5E57DrXVuCFpceGGMQLAiwIQcrDvL6yo3CFAubwqxGR0ZjgqhtjnAyw0IvFBGIAgQYEIOgk3kJTACjAAjwAjIHwEmZPnvIa+AEWAEGAFGIAgQYEIOgk3kJTACjAAjwAjIHwEmZPnvIa+AEWAEGAFGIAgQYEIOgk3kJTACjAAjwAjIHwEmZPnvIa+AEWAEGAFGIAgQYEIOgk3kJQQvAvkGA5QKJZTK0iNv5ebmISklHbGVojhKV/A+DryyIEeACTnIN5iXJ18EsrJz8eSQ6Rjc5zF0eri5w4VQmMyVm/di+cb3RX3FqHAsmzvakgFKvqtnyRmB8ocAE3L523NesQwQWLBqOzZu+1BIOm/SkBIJ+fTPf6LPiDnYsnQiGtSrjSXrd+PA0a9wZPtCp1q1DGBgERmBcoUAE3K52m5erFwQSE5JR3ZuLnoPm4UXB/cskZDfWLUDv/11DusWvCSWFn89GQ8+MRq71s6wSc0ol3WznIxAeUaACbk87z6vXfIItO/1EkYO6F4iIY+buRLRkWGYNOoZy1pub90PK14dg1bNG0p+fSwgI8AIFCHAhMxPAyMgYQScEfLglxbg1jo3YuzQnpZVNOkwFNPH9cOjbZtJeGUsGiPACBRHgAmZnwlGQMIIOCNk0pDJkGviC31YQ5bwPrJojIArCDAhu4ISt2EEAoSAI0IuKDACCkCpUIDukM+cPY8188cJCfkOOUAbxdMyAl5AgAnZCyDyEIyAtxEg/2NjgRGd+r6CoX07o9NDzaHRqEFuTjl5BWI6nUaJH375q9DKehIa1K+Nxet24eDRk2xl7e0N4fEYAT8gwITsB5B5CkbAXQRenL4Ch49/bdNt/+ZXcdMNcbiSlC2+j4s25T5etvF9rNq8V/x/qD4Ea+aPRaM76ro7JbdnBBiBACPAhBzgDeDpGQFXESgwGpGelS/+oxKmV4v/6Og6OycXiUmpiIuNYf9jVwHldoyAxBBgQpbYhrA4jEBJCNDdsVk7NrchLdlZWE1GlBFgBOSBABOyPPaJpSznCBAZJ2fkIjvXdH9sLiFaJaIqaJmUy/nzwcsPDgSYkINjH3kVQY5AXn4BrqXkOFxl5UgdNGplkCPAy2MEgh8BJuTg32NeocwRIO34emoO8g1GhytRqxSoFKFjLVnm+8ziMwJMyPwMMAISRoDcnLJyDUhOzytVyqgwDfRaFadelPBesmiMgDMEmJCdIcT1jEAAESBCvpxocnNyVqpWDGFCdgYS1zMCEkaACVnCm8OilW8E6Kg6LSsPGdkGl4CoEKJCuF7DR9cuocWNGAHpIcCELL09YYkYAYGAocCIq4VBQFyFpEp0CFRKhavNuR0jwAhICAEmZAltBovCCJgRIO04KT3XEibTVWQonGZ0GLtBuYoXt2MEpIQAE7KUdoNlYQQKEcjJMyAhNdcjPGIitNBpVB715U6MgK8R2PrBUXzzw+9YOH14qVOlpGXg5He/4PzFeNS+sRqaN75NhIYN5sKEHMy7y2uTJQIUIvNaco44svak0JF15SidCKnJhRGQGgILV+/A4ePf4PDW+SWKdvVaEvqPeQ3nLlzFLbVr4I+/L6B+3ZpYv/BlRIZXkNqSvCYPE7LXoOSBGIGyI0BknJWdj5RMU7xqT0tkqBr6EFOcay6MgJQQcIWQF6zaju17PsEHG2ejelwlnDn7H7oPnIKxQ3tiwFMdpbQcr8rChOxVOHkwRqBsCLjj5uRsJnaDcoYQ1wcCASLk9z88gYG9H8XuA5/h7LlLaHNfI0wf1x8x0RFCpJZdR6J1i0aY9fIAi4jPT1iEawnJ2LV2RiDE9sucTMh+gZknYQScI0CGXKmZecjMcc3NydmIoToVIkLZDcoZTlzvXwSIkNdvPYiaNaqgW4eWgpD3ffQl7mtyB9bMH4fc3Dw0ajcI44f3Qt8e7S3Cvbl2F97ZfQTffLjKvwL7cTYmZD+CzVMxAqUhkG8oQHyy43jVniIXG6WDWsVxrj3Fj/t5HwGzhnxs5yJoNGoxwfKN72PFW3twZMdCqJRKPPjEaMwY1x9PdGplEWDN2/uweN17+OHIemjUwWm0yITs/eeNR2QE3EaAtOPEtFzk5ttmc3J7oGIdtGolKoazG1RZceT+3kPA0R3yiVM/Yej4hdiydBJq3RiH+7uMxLSx/dDzsdaWiYmwibh/OroBqiB9yWRC9t5zxiMxAh4hQPfGOXkFgpB9UYiQyT9ZwQZevoCXx3QTAUeE/OlXP2LYK4vw9rJJuLN+HdzZdgDGDO6B53o/ahl93vKt2P/xlzjxwVI3Z5RPcyZk+ewVSxqkCBAhU0QuD72cnKJCgbsoghcTslOouIEfEHBEyK8v34q3dh4WZFsxKhzdBkxGtbhKWD53tEWiJwZNE3V0zxyshQk5WHeW1yULBMjNKTM7H6lldHNyttiIUDVC2Q3KGUxc7wcEiJC3fnBMBAapUbUSjn7+PVZv2YcHW9yF16cMFRK8+/5RzFm8BVPH9EXTRvVx5MR3IKOuhdOHoX3rpn6QMjBTMCEHBneelREQCHjTzckZpOwG5QwhrvcHAovW7MS6dw+IqFuZWaZMZg80a4jXJg5GZIQp6Ee+wYB5y94VxGwuw57tgmH9ugb1SQ8Tsj+eQJ6DEXCAABlypWTkiXzH/iiULzmyArtB+QNrnsM5Anl5+bh0NQFhFfQW/+PivbKyc3ElPkEEB9FqNc4HlXkLJmSZbyCLL18E8vILcC3Fu25OztCoHKmDRs1uUM5w4npGIBAIMCEHAnWes9wjQHfHCSk5yDN4Fq/aUwA1KgViIjnOtaf4cT9GwJcIMCH7El0emxFwgADdG2fnFSDJR25OzkCPDtcihN2gnMHE9YyA3xFgQvY75DxheUeACPlKYjb8qxsXoU7pJuIqshtUeX8Oef3SQ4AJWXp7whIFMQJ0VJ2emY/07LJlcyorRGEhaoSFcjaosuLI/RkBbyLAhOxNNHksRsAJAmRZfSXJ5OoR6BIXHQIlRQ3hwggwApJAgAlZEtvAQpQHBIiMk9Nzxf2xFArdI0eFcZxrKewFy8AIEAJMyPwcMAJ+QiA334DrKb6JV+3pEipFaqEN0sw5nmLC/XyHQEp2bVBIdbKfUEAhAuP443OE7m/fLcqLIzMhexFMHooRKAkBuju+npyDfF8FrPYQerVSgUpR7AblIXzczU0EUjNrm3oIFrYya/Tx54hQJmQ3t4qbMwLBiQBpAVk5BiRn5ElygVEVNNDrVEEdklCSwJdDoVLTa5s0Yxj9+m942FlZoM0asiy2iYWUOwKXErIkvQSOcy3p7Qka4dJS6wRkLeERTMgBAZ4nZQSkhAAZcqVm5SEz2z/xqj1de4UQFcL1HOfaU/y4n2sIpCfVocvjwkvkojtk06Vy4TG2D+rDohwTMv190gmWSiWNcLKsIbv2HHGrIEcgLT1TZJiJjgx3aaWUpSYvz2DJTlNSJ0OBKdexHArlTFY5cYO6npiCCqF66EO0TpdEiQGSklMRFxvD7lVO0SofDTIS6lgMukysXGTg5cvPYTH2hExEPP2NTUKGGeP6l7oB85Zvxeadh23aNLqjLt5eNsmrG8eE7FU4eTC5IUDEOn72ahz74rQQ/c7b6mDp7BdQqWKkw6VcvZaE2W9uxsnvfxP19W6+ERNfeBr169a0a09v30npuciRiJuTs73RaZSILsEN6vzFqxg6fiHOXbgqhune8QFMffFZaEqw0B45abEFU0oq3/WRlhg7tKczEbg+yBHIiK8jbBVM1tX++ze08l82yB4+/jVmv7kFiclpeKJTK6eE/Nqyd/HfpXi8PKyXZRydToO4yhW9umNMyF6FkweTGwKUl3XnvuPYsnSS0Pqen7AItW6silkvD3C4lJdnrUJyajqWzx0NhVKBGW+8hWsJSVg1b6xNe/rByc0vQEKqtNycnO1PTIQWOo3KrtnglxaINHlzJgwS6fB6Dpkhksc/1q6FwyGXbXgf7Vo3wY3VY3Hyu18xfOKb2LZyKhrUL7SydSYI1wclAplXbg7IukLjbAk5MysHqekZoNzMITqtS4RMf/eUs9mXhQnZl+jy2JJH4IlB09C+dRMMerqTkJXenF+cvgI/f7LRodVxnxFzULNGFcyZ8Jxo//6HJ7B0w24c27nIZq3k5hSfnAPSkuVU6Mi6cjE3qJS0DLR4bLg4nqNjOipzFm/BlfhELJ0zyqXltekxBk91aYPBfR5zqT03Ck4EMi/e7Be/4+J+zvpqtoRsRnfmos0wGAwuEfJHn36DZnffJq612tx/N+658xavbxITstch5QHlhECTDkMxe/xAQcpUfv3jX/QYPB1f7luOyPAKdks59vn3GDl5Cdq2vBvdOrTE/BXbMOCpjuLYy1yIjDOz85GaGdh41Z7uQ2SoGvqQojjXZ/+9iM79JuH4e2+ickyUGHbLro+w5/AX2LV2htNp6Ji7Y5/xWPHqGLRq3tBpe24QvAhk/VeoIfvY77i4n7P+hrIR8r6PvsS/F65Ap9Xg5zP/4OiJ77Fw+jC0b93Uq5vFhOxVOHkwOSFAx8p3PNjfhijM5HNk+xuoWiXGbjkXr1zHoHHzcUvtG/DFNz8jRKfBxkUTcHOt6pa2NO7lRHkYcpW0X9ZuUKd//hN0MmD9krJj33Gs2rzH7mSg+HgZmdnoM2I2wiqEYtObEyRjzSqn5zSYZM0+R4Rc5IdsNukq/q/ZT9lb9SE1/3QIo6sacvHOE+auQXJKmt1VVVn3igm5rAhyf1kjQBoyHT+3a9XYJQ35ySEz0KrFXRj2bBeQZfa0BZtw4tRP+Gr/cqhVKnFEnZNfgPx8acSr9nRz1GoldGqlsI42v6R8unuxxdjNFQ2ZrKxHTVkijrY3L5mIqMgwT8XhfkGCQM4/dQutrEuiWt98H1LLu4T85tpd+O6nP7Bl6USv7gwTslfh5MHkhgDdIT/yYFM81/tRIXppd8ik7TXtOFRYYdMdEpVfzvyLnkOm44ONs1G3Vg1ByGTMlSdzQtZqVNCoFVAqFHB0hzxr0WbEX08q8Q45NT0TL0xegqysHKx+fSyTsdz+MHwkb85fdYvcjQutrM1+yUVW16aomubP3qjX1bElZIOhAAUFBZi9eAvy8w2YPrYfVCqVeAGlF+3+Y+ZhYK+O6NDmXoEEGX91btcCN9aIw5mz59F/9DzxmzHkGe/aRDAh++jB42HlgcDad/Zj1/5PhZV1qF4nXHusraw37Tgk7ovMb8Lte72EWjfGYd7koQgN0YHelD/58jT2vjUXSoUSiem5iAnX4XKitCNzlbY79KMUG6kTmanM2aCeGzcfEWEVxGmCIytrMoSrFheDcUOfBFmwPjV0hvDrXjRjhLDOpqJUKlE11rtuIvJ4ylhKMwK5f5iMAouig5hrzNFAfPNZe4stIe/Y+wlmLHzLZmPIs4Lc+VJSM9Ci83BMHv0MenVtK9rQyRjdHZtL10fux5QxfYWFtjcLE7I30eSxZIcAab3jZq7EZyd/FLLfcWstofXFVjIZL5HRFt2XfvPhKvH5tz/PYeXmPYKkQ/UhaNzwVnF8fUe9WsLfODEtV0S8gsKINJkadREJiwhjmXmoGE5uUEr8+98V8bJy4fI1gQP9IJFWodGoxeduAyaLF5mF04eDfLXJqrp4IX/kEx8sld0zwgJ7D4Hc30hD9p//sdnfWVPvjzIvgjTnpJQ0VI6JdikwjicTMiF7ghr3CToE6Fg2Ly+/xIAgxRdMRE5HXZERJkts+sO/kpRtSWBDUa+up+SAInXJqWjVSkHCV5NzLKnx4qJDLC5gRLak8VYIDZHTslhWiSCQ94v3XYVcWZrm9rITsivzlLUNE3JZEeT+5R4BcnPKyMpHWlaRm1OIVoVQnUpozHIqlSJ1wmUrM6co9na4Xo0K+iI3KDmth2WVFgJ5P90i7pBNxSqWtY8/qxswIUvrSWBpGAEfIUCEfMWBmxNFvUrPMiAnT9qJJcyw6LUqQbyk2RcvcRVDhIEXF0agLAgYfrilMI9EIRlb8kz49rP6Libksuwb92UEZIEA3bVSnuPsXHvSVavo+FcjInZJvRDXxkaFICktV1iJFy+k8VPeZDL44sIIeIqA4ftbLHmQzYZdJeVH9ma98u4znors1358ZO1XuHmyYEOA3JuuOdAozeuMrKBBvsGIjGxpR+0KD9VArVIIQi6pVI7UQaOWRpq6YHuOyst6Cr69tWippviWJRcv1isbMyGXl2eM11lOESDt+HpqjiDc0kq1GD0uJUjXDYriV8dG6YQmX5oRGhF2pQgda8nl9Hn3xrINp241GQiaHI0tWZ98/VnZ9HdviO/zMVhD9jnEPEEwIkBW1XRMnZSe53R5FULUQvtMyXDe1ulgPmgQHa4VLxVpmc7liw7TgI6vxY8qF0bATQQKvqrvoIcXVWGH8hihbM6E7OZWcXNGQD4IuBuvmu5nE9Oca9P+RoDcnIiQ45OLXLacyWAd59pZW65nBKwRMHxRv8gP2SqmtTl2tfjXrDl7sV5536+y2AjWkGWxTSyklBAgq+r0zHyku3EvTDmGw/QqyeVHJjcnctnKcmCUVhLmYSFqhIWyG5SUnkm5yFJw4naLqMX1Yl9+Vrb8RRYQMSHLYptYSCkhQHfHFATE3UIBN8i/15FFtrtjeaN9aIha+Eo7cnNyNj4FC2GLa2cocX1xBAzHby8hUpe4UrbcKZsjbBX9W7Z6ZaufZbEZTMiy2CYWUioIEBknpeeKMJnuFjKeIo30qgdk7u5cztrTsSAZcpXk5uSsP4XTjA7TMik7A4rrbRAwfNJApHMS5Guu8cNnVZv/yWInmJBlsU0spFQQyM0z4Hqq59G3IkLVKDAqkJ7l3IDKl2uOCNWAXIrJh9rTUilCC8oKxYURcBUBw5EGAYllrWz7k6siBrQdE3JA4efJ5YQA3R1fc+Ia5Mp6qlbU40oU+EfUAAAgAElEQVRSliXutSt9vNmGApYQmcan5IgkEp4W0vgrR+k4gpenAJbDfoaPGgZk1ap2puQxUi9MyFLfIZZPEggQGWfnGMqkUZoXQne3WpXCK2N5Ak5MhE6E80y3ir3tyTjUh6J36XXsBuUpfuWtn+EQEbLJZc58Z2z5LMJoFpl2ebNe9cgPsoCaCVkW28RCBhoBd92cnMlLUa/ouJgiffmz0N0vRQ+j6GIUm8Ebhd2gvIFi+RjDcLBRQBaq6ng6IPO6OykTsruIcftyh4A5N7B1BqSygqDVKBGh14hIX/4sZMhFAUo8MUorSU6y1BZ30hzn2p9bKcu5DPvuLgxRbbbkKh6xq1CBLozkZR/By7N6VafvZYEXE7IstomFDCQCBoMRV5Pdd3NyJjNZKWfnGZBllerQWZ+y1FPEMIqyleCDl4AqUSFQqTh6V1n2pzz0Ney9x2ReTY+KH/9VdflOFvAyIctim1jIQCFAsZ09dQ1yJjOlM4yN1jlM3eisr7v1pL2Sdnw9JRf5Bu8fk1PEL/KzZi3Z3Z0pX+3z37+n0Mq6uF+xbz+run4rC6CZkGWxTSxkIBCge+PcvAIklJIBqaxyUZYlOpZL84KBVWmykPEVKSS+jKdN+Z+JmDnOdVmfiuDtb3iviZXZlmmdvozQZR5f/fg3sgCVCVkW28RCBgIBsqyOT8pGGTyDXBKbol6V1QWptInIzYnIkrI50UuGr4pSacqpTJo/F0bAEQL5O5sWnlYbC/Mim0+vfftZ1eNrWWwIE7IstomF9DcCRMaZ2flIzfR9HmNyGwrRqEQEMF8UIuOs3AKxHl8XCnxCbl1Myr5GWp7jG7bfWxShq5SrZPPqSrpqdrde/eQpWQDGhCyLbWIh/Y2At92cnMlPeYZTs/LEEbk3i16rQnioWmjH/irsBuUvpOU3T/679wbmDrnXSVmAxYQsi21iIf2JALk50V2rOxmQyiqfRq0UQTbIP9hbhU6O6Qg52cPY257KQRp/JLtBeQpfUPfLf6e5rXW1u6pucZXZxf7qPl/JAlcmZFlsEwvpTwTICtmfGqV5bUTIuQbTUbk3SpheLYysEn1olFaSnGTRTXfXXBgBawTyN7cojNBV8nm1KUKXd+vVz3wpi41gQpbFNrGQ/kKAtOOEtBzk5fvO+KmktdAPUVy0HpcTs8q8XOHmFKkTGje5bvm7aNQKxITr2A3K38BLfL78TfcFREJ1vy8CMq+7kzIhu4sYtw9aBOjemCJYBUKjNIMapqcsTMYyG5NFhWlF4ojUTM+zOZV1o8kvmUJ1shtUWZEMnv55G+4vigmiUJjyH5tdn3z4WT3gc1mAyIQsi21iIf2BAP04XEnK9lqMZ09lrhIdgutl0GzNQTqu+tjNydn6TBp/CBOyM6DKUX3eupaFfsdmGjZfKfv2s/a5E7JAmQlZFtvEQvoaAXJzysjK93mADlfWQeEtKT60p5p6pUiduIf2ZuxtV+R21CZcr0YFPbtBeYpfsPXLXfOAOSSXX//VDvpUFlAyIctim1hIXyNAx7ukHUulkO9wepZBpEl0p5CbExEgadhSKXEVOViIVPYi0HLkrG7t1xjW5vNw3dDjgV66S/MzIbsEEzcKZgSIjJMzcpGd610f4LJgplYpUDFch3g3klrQXS1ZN/sq9ran6wnRkksXx7n2FL9g6pe94sFCP2SjX//VPX9MFjAyIctim1hIXyJAOYm96f/rLVkppSEdpae7GOea2lPGJSJkqRXK/0y+1lzKNwLZy9uafJqsw6taPhf6OvmgPmQEE3L5fvJ49bJAgLRjykmcb/C/a5ArAFWL0eNSgnM3KJVSASK9QLk5OVsLafwUjYyzQTlDKrjrs5a0tWjGFodks+Nx4b900iNirhf73vzZk3r9yCOyAJY1ZFlsEwvpCwToj56icSWnB841yNm6KIcxkZmzLE3R4VqQpu+qNu1sXl/UR4VpQHfc7AblC3TlMWbW4ocdxrI2S19ka236xlufQ0d9LAuAmJBlsU0spC8Q8He8ak/XQPfCiWl5JeYx1mlUILKj+2YfJnPyVHybfhzn2iswynaQjIXtTBoyCrM7kR+yHz6HjjksC8yYkGWxTeVbyOuJKagQqoc+ROsyEBmZ2UhNz0SVStEOj0npqDotKw8Z2e5ZMbssgBcbEuGG6VVISHV8N0yEnZaZ79fY254uLyxEjbDQ0t2gcnPzkJSSjthKUaxNewq0RPtlvPGInzIgW6cANaLCWCZkiT4SLJZcEDh/8SqGjl+IcxeuCpG7d3wAU198Fhq1qsQlfPrVj5i3/F1Ln/c3zMYttWvYtadwklcl5ObkbE/oSDo7x2BHupTqUK9VlkjWzsYNRD0FPqE77+KFTixWbt6L5RvfF1UVo8KxbO5oNLytjlMxad+HvbIIK14dg1bNGzptzw0Cg0D6/EcK8yCb8x/7598KL30YmAW7OStryG4Cxs39h8DglxYgrIIecyYMwpX4BPQcMgNTx/TFY+1aOBTi+Jc/YPjENzHo6U7o0v4+REeGQ6fT2mnWpB1T7mEKkymXQgRGAT+sXyLMbk6kOVNCDLkUCqcZHWbvBnX65z/RZ8QcbFk6EQ3q1caS9btx4OhXOLJ9YanGYGfO/if6ZWZlMyFL/CFIn9exMFKXf92Rw8cflDgyJvGYkGWxTeVPyJS0DLR4bDjeXjYJje6oKwCYs3gLrsQnYumcUQ61q+4Dp+DWm2/EaxMHlwoYBdso6fhXykhTXmMYFeKonUpkBY0weknOkK5RWkl4VorQQquxPel4Y9UO/PbXOaxb8JLoFn89GQ8+MRq71s5A/bo1HQ51LSEZTw6dgRcH98SMhW9hwdTnWUOW8EOc+uqjQkMWdGyJXe37z+Gv7JcwKkWiMSHLYpvKn5Bn/72Izv0m4fh7b6JyTJQAYMuuj7Dn8BfiB7p4SUxOQ8uuI9HmvkbIy89HRmYOmt9zGwb06ogQXdHdM/n1XksOTAYkb+xi1Yp6XEnKgkqpRKVIrUgTSRq/3Ipw04rSQWnljzpu5kpER4Zh0qhnLMu5vXW/ErXerOxc9Bv1KlreeydGDOiGJh2GMiFL/EFInfNYQO6QIyYxIUv80WDxpIyA+fjyy33LERleQYi6Y99xrNq8B8d2LrIT/bc/z+GJQdPQo1NrtGhyB1LTMjBv+VY82rYZpo/rJ9oTGWdl5yMl0zv5hgOBX6hODa1GIQg5O9eADC/lTg7EWkjD1+tUFlKmK4pb69yIsUN7WsQhkqX9o320LvQSQgROhbRi8m9mQg7ELro3Z8qsziW4F5sjd5Xkfly2+ojJe90TNECtWUMOEPA8bekImDXkT3cvRqWKkU41ZDMhn/hgqTAGorL74Gd4dem7+PrgSkswAin76br6TJBvMp36UTIMuZcwvdpiSU0ES3s38YU+TjVk83H2E51aoYI+RLR/a+dhtG5xFzq3uw/tWzeROzRBKX/KzC6my+OSSnHH4+LtPKyPnLZHFngyIctim8qfkI7ukGct2oz460kO75DN7beumII7C61yd+z9RNwr/u/YRgEgGUGlF96/yhlRSh5BP2py1o4JfzquJitxsq4mDZfukM+cPY8188eJ7SntDpkMuN5+zzbYw+J176HTw83R6aHm4hibi/QQSJ7W1a9ZnsxhOqOmmyz3pV6YkKW+Q+VYvufGzUdEWAXMmfCcnZV1Wnom+o+Zh4G9OqJDm3sFSkPHvyHuU9+cOQLXE1Px0syVqFolRnw2GAqQkVMApcKIVBkfWZNGSX7JhoIC5OYZkZkjXy2ZXLkMBqNw21KplCiysp6EBvVrY/G6XTh49KTFynrTjkM4euJ7YYXtqPCRtfR/LJKmdhexrE3hL82huHz/OXoWE7L0nw6WUNII/HP+svBDvnD5mpCz6yP3Y/rYftBo1EhJzUCLzsMxefQz6NW1raindqOnLgMdX1O5t1F9vD5lKCpGRYjcwrn5BSAfWEpNSH7IciukRVIQkOspueLFIzZahyuJ0kkZ6Q6eWrUSRMhklKYRma204m5x2cb3sWqz6b4vVB+CNfPHWqzs56/YJuwIvvlwFROyO2BLqG3ipMdLClHt0++jZ78nIRRKFoU1ZFlsU/kW8uq1JOGPXCHUdFforNBRp1qtEveRdBxK/sZEyFRCtCqE6lSWz87GklJ9VAWNuH4zx7UOD9UIbSNNhnfJ5FNNd+AUS5wKETL5J5PmlJ2Ti8SkVMTFxnAyCik9gF6QJXHiE14Yxf0hKs7d5X6nAPRgQg4A6Dyl/xAgQqZgGtYKcUy4Fuk5BuQUkoH/pPF8JrVKiZgIk0YpMuEUlrjoEMSnyMv1iV6I6O6YTirMhQJ3xUaH2LhBeY4W95QqAtcn9CjyQ4ZVTGtzbOvCf81hQywxr8tYH/PaDqlCYiMXE7IstomF9AQBcnPKzM63uzOm7EkVw3UiGYNcCpFxVm6BWI91IbehEI1KRB6TQyENuIpIlmG6QrAuEaFqQdTWvslyWBPL6DoC119+0hIUxGxtbSbdolHMqRdN33ijvvLrTMiu7xK3ZAR8gEBp2ZwiQjXCL1kOblCUspCidJF27KjEROiE9bgcQoES7nQXnlzCCwRng/LBH4KEhowf91ThXbHZkKv4v2YjbO/WV56/TUIolCwKa8iy2CYW0l0EyOiJ7lrNd5SO+leL0eNSQpa7Q/u1PRk6xUaFiJzNFPLTUdGolaD75WtWR8B+FdLFyUR0rkhdqUfs9PJBAUOItLkEHwLxY3tZFmUxsrbKe2ybo6koH7LZ/djT+ipvbJUFmEzIstgmFtJdBPLyC5wSFAXYoONrs5GUu3P4oz25ORHhJhUapZU0J5FYvsEoad9kMtyiY2pnpxIUUlOjUvoDXp7DzwhcGdO7KNuTVSxrcSztw89VFr3j55V6Nh0Tsme4cS8JI0BH0QkpOcgzOHdtIjeixLQ8SWZLMmuUpPk6c9MizSEuRo/LEtX4yYI6KoyM0rJN/qelFHKDomN41pIl/EfmoWiXR1EUtiLyNevAZjL21eeqi5mQPdwy7sYIeI4A3Rtn5xU41SjNM1CQjTC9SpLZn8yBM1IzXcvmRNo0GUS52t5zlN3vSS8+FJCF4m+7UqzdoFxpz23kgcClkX2KNGGzRuyHf6su2SILgFhDlsU2sZCuIkCETMEynOvGRSMS8WXnGEq9b3Z1fm+1sw6cYe3m5Gz8KlEhuJ6WIyJgSaWQ5TTdDSekOjZKcyQn3Z2TSxdpTlyCB4GLI/oGZDHVl20OyLzuTsqE7C5i3F6yCAir6cx8pLuZAYmOhilQBfkrS6UUD5zhqlxSC3xChEraMeWfzjfYujk5WxNp/Gat31lbrpcHAheGPWsVkYvCZxrN2ZGtNGcR78bqTrnQK9miSbtfX33FW7IAiAlZFtvEQrqCAFlWX/GQVMmtCEYF0iSQfMJR4AxX1m9uQ8e9mTkGl4+H3Rnb3bZkbEY6bnKGa8fuxccnLZnvkt1FXbrt/3u+n5ldC/Mim9MjKwpZ1upzkaOyVUIKz+pvWLVJuqBYScaELIttYiGdIUBkTL6tdH/saSEf2CtJtpGwPB3L036lBc5wdUypBD4xRxcjozTaH09KiJZcurRMyp6AJ8E+5wf3L/EO2aw6W6yti90tl6X+htUbJIiGvUhMyLLYJhbSGQK5+QaRdKEsJVSnhlZDQSs80+bKMre5r7PAGa7OIYXAJ2QpTX7gxaOLuboGczvyXSbXLy7yR+DcoIEBWUTNtesDMq+7kzIhu4sYt5ccAnR3THGRyQ+3rIV+/Ol4lfyY/V1cCZzhjkyBDHxCbk50XF1SdDF31kEafyV2g3IHMsm2/Xfgc46jfZQUJcRL39+0bp1kMbEWjAlZFtvEQpaEABmFZOUYPL6jLD4uWTdHVNDYJD7wF/quBs5wVZ5ABT4hw2h6saGAK94K5xkVphGW2mx17eruS7PdPwMGmfNGFAloJl3zNz74XGvjWmkCUkwqJmRZbBMLWRohX/ZyTmD68c/NMyIzxzaRgy93wZ3AGe7IQVGvktPykOemhbM7cxRvS5bR5N/tjpuTK/NxnGtXUJJ2m7/7DRaBMosSRpjY19efa29aLW1gCqVjQpbFNrGQjhAgQyGyis7Idi3YhKsoUnCN2Gid8Gf2V3E3cIarchHRh+k1XifHkuYni+jYSB2ue+Dm5GxNFUJUCNdznGtnOEm5/mzfoVbi+UAVFjb99qp2nc2rpAyLRTYmZFlsEwvpCAEKJ+kr3+HwUI1ww0jL8r2W7EngDHeeCH8GPqEkF2RQ7atoYVWiQ0B37VzkicCffYaarh0KuVj4IQtratssT96uv3nLSlkAxoQsi21iIYsjQNox5QD21h2lI4TJBza+DC47ruwaaZR03+pJ4AxXxqc2/gp84ml0MVfXQe1I448OYzcodzCTUts/nx5mJ05xPbl4A2/U3/LOCinBUKIsTMiy2CYWsjgClIqQSMyXRa9TIUSjEsTvq0IaJRVPA2e4Kle4Xi1UkDQX42K7Oq51O4ouRgFJyurm5GzumAituKPmIj8EzvQaXpTtyeru2HyH7Kt/b9m6TBZgMSHLYptYSGsEyM3pWrLzDEjeQI18adOzvGctbC0TBc6oFEkZkDwPnOHOGuMqhiA+KQeEn7cLWUBX0Kv9Yp0u3MOidCKRBhd5IXDmqZGF4TILDbl8mHLROqVjve1MyPJ6UlhaWSBAZJKVnY+UTN/f7RIglJc3KlwjXgC8XYjsKftRhpuxtz2Vw1eBT4gXY6NCRIYtynfsjxIZqoY+xJTdiot8EPitx8jCy+KiO2NTBC7ffq6/Y4lDkOjqi+6xVRLJv80asnyeZZaUwtgajfC2m5MzYCnABQUd8SZxejNwhjP5revpWDk1I8+rxEkGcBS8gwjZn4XdoPyJtnfm+vWJUd4ZyM1Rbtu12K4H/ZZMf8MU43rGuP5ujuib5kzIvsGVR/UBAvQ2S9a7dE/pz0I6WFyMHpcTsrwyrS8CZ7gqmLcDn9DxMbls0bE7Wb37s1ASDnOoUX/Oy3N5jsDP3UcHJB/y7e8tshH68PGvMfvNLUhMTsMTnVoxIXu+pdyzvCJA6fu8EYrRE/zMaQC94c7jq8AZrq7Lm4FPyKWKTg98aSxW2rroZYDu4rnIA4H/dXsxIII2eH+hzbyZWTlITc/AojU7EaLTMiEHZFd4UtkiQNpxQlpuQGJMm0Gje9LEtLLFzPZl4AxXN5f8PuOidWU++i9yc8oWfqSBKCQDhRzlFI2BQN/9OX/sMjYgGvKdHyxwKOzMRZthMBiYkN3fSu5RXhGgux7yN0708x1lcbxDtCrQMWlZ5PB14AxXnxGKeAUFabaeG8fRfXRGVr7I6BTIQoRMd/Ic5zqQu+Da3D90fqnIgMvcxRwVxPLZYvdl+sYL9XftY0J2bYe4FSPgBAEi5CtJgdPCrMWjH3+6wybraHeLPwJnuCMTRb2iLFme3P1SdDF6OaH+gS4UuIvWwoQc6J1wPv/pTi+L2NWmUF1FxTb4h/frG+1/nTVk59vDLRiB0hEgNyfSwvwRwtKVvSBr4orhZMTkfpxrfwXOcGUd1IZ8h0N0Kreto4n46O7Wn25OztYUEaoGvSSwG5QzpAJb//2jE0x+yGb/45L8kL1cf8/BeUzIgd16nj0YECBC9meSB1cwI8tekivdjTjX/gyc4coazG0o6lV6lgEU+czVIiybFb6PLuaqPOZ27AblLmL+b//tIxNM59FmPdnqOFroxT763PjDV20WazAUoKCgALMXb0F+vgHTx/aDSqUKuC0Cuz35/5nkGV1EgAy5KKSkJ8fDLk7hcbOqMXpcScgqdvDmeLhABM5wdWFkoVwxXOOy9bqIkhWp83mMb1flt25HLz3kM84GXp6g558+37Sf6J+Jis3S5PBcm2927P0EMxa+ZfPdrJcHoHvHBwIin3lSJuSAws+Tl4ZAXn4BrkngjtKRjBVC1CIYRkpGntNNDFTgDKeCFTZwJ/AJ3aFTNC53TgdclcMb7ehlQaNmNyhvYOmLMU49PCkgVtZNP5rti+V4fUwmZK9DygN6AwHSjq+nls3FyBtylDYGxVNOTstDnqHkcJGBDJzhzvqrxehxyUngE3N0MXpJCpSbk7M1aVQKxERynGtnOAWq/uRDUwIydbMjswIyr7uTMiG7ixi39zkCdI9Ex9RJ6c61T58LU8oERFBheg0SUku2NA504AxX8XEl8AkZctGJgC9TXroqb2ntCPMQdoPyBpReH+PLNlMDoiE3PzrD62vxxYBMyL5AtZyMmZubh6SUdMRWivKqy0kg4lV7umWUmzc7z4AsB+E8pRA4w511lRb4hCyY6Y62tJcPd+byZVsR6rRi6W5QdAITn5CEShUjoVa5nsoxISlViB4THeHLJQTt2F88OF24PClE6sXCYra6tqza+/X3HWdCDtqHqrwvjAhz5ea9WL7xfQFFxahwLJs7Gg1vq+MUmk+/+hHDXlmEFa+OQavmDe3aC+vlzHyk+ykDklOBnTQwR94iP+niRSqBM1xdo06rQphOJSKiWRdaI93NUv5pCl8qhxIWokZYqGM3KHoGx81cicws055NG9sPPR9rXeKyiLzXbz2AzTsPi9jHofoQfPPhKjnAIDkZP281w5Ebsq2c9m7IZa6//9NpksPCkUCsIctim6Ql5Omf/0SfEXOwZelENKhXG0vW78aBo1/hyPaFpVq4njn7n+hHP4QlEnKBKQiInEq4Xi1cOazjOUspcIY7WDoKfELRxUibccWAzZ25fN02LjrE7nnMys7FA91ewIgB3fB094dw/MsfMGrKUhzeOh81qlZ2KNIbq3bgg0MnMLRvF3Rocy9y8/IQV7mir8UPyvFPPDCrMPKWOQAXpV6k3Iu+/dzyxFRZ4MmELIttkpaQ9AP121/nsG7BS0Kw+OvJePCJ0di1dgbq163pUNhrCcl4cugMvDi4p3A3WDD1eTsNmTSRpPRcyd9ROlogHZHGJ+UI/2QpBs5w9QlSqRSoFKHD1cKXInKLIl9lSuohfERlVOgeOSrMNs61+YTm9EdrodVqxGo69hkvyPnp7g/brY6e29aPj8bs8QPRrUNLGa1emqJ+et8s0x2y+djaclxdeEzto88PfD5ZmoAUk4oJWRbbJC0h6bgvOjIMk0Y9YxHs9tb9StR6SSvpN+pVtLz3TqGZNOkw1CEh5+YZcD3Vvzl1vYVsqE4NrUaB5PQ8U0pACQbOcHWt1oFPYiJ0IlZ1pkyuEIqvsVKkFlp10R3xjn3HsWn7hzj4dlHkppGTFuOmG6pi7NCedhAdPfE9XpiyBE91aYM//r4AnU6Dzu1aoHO7+1yFk9tZIXC8xZyA4NH6y0kBmdfdSZmQ3UWM22PwSwtwa50bbX7AiGSnj+uHR9s2s0GItF4icCqkFdN9pCNCJs3yWgBy6npzO8WdcXY+IkM1kgyc4c5aq1bUIzk9F+GhapeDhrgzvr/aqpUKVIoqcoNa9+4BHPrka3GaYy70fIaF6sXzW7y8s/sI5i55W7xI3lr7Bpz5+z8s2/A+Xp8y1O5Z99ea5DzPsWZzA2Jl/eBXr8gCNiZkWWyTtISkHzAy5Jr4Qh+nGrL5OJuSgFfQh4j2b+08jNYt7hJaRvvWTcRRaJ7BiBwPEjZICRnyOdaHqEGafm6ePIyfSsKPgmvoNCpk5ebDYJDXUXXxNVGWLgriQkel7mrIRMjb9xzD3reKIj1NmLsG2dm5eHPmCCk9frKQ5ei9rwVEzranJgRkXncnZUJ2FzFuD7pDPnP2PNbMHyfQKO0OmQy43n7vYxvUFq97D50ebo5ODzUXx9iCkPMLZHl3bL0wun/Va9UiLjStR85Fq1FCS4Scne9RNigprZ0I2Ry9y3yH/MPH66DRqIWY7Xu9hL492jm8Q7a0P7IemsKjb3ohzcrOwfK5o6W0TFnI8nGTeSKUtTlkdVGiicKUi2Tf5YP6h74eLwt8mJBlsU3SErLIynoSGtSvjcXrduHg0ZMWK+tvfvgd85ZvxRvThqFmjSp2wlsfWdORdmpmngiwkZgm7chcznaBAmdQfuGoMA0uJ8rLUtx6bWZXrmTaF51aREyTazH7gqdn5YFyQGfn5KJJhyEYP7wXejuwsk5Lz0T/MfMwsFdHYVGdmp6Jtj1exLM92uP5Z7vg5zP/oPewWcJ+one3tnKFJWByf9T4dWGxb0ovYSr++Nz+W0r7KP3ChCz9PZKchPRWu2zj+1i1ea+Qjfwy18wfi0Z31BWfP/nyNEZMXIzd62fh1jo3lErIdBx6NTlbHI+G6SnwhEyNuqwCZ9APPxRGQc5yLGSZbH5RKi3wiRzWZu0LTjmT6Vrh2BenQYZc5jJ59DPo1dVErimpGWjReTisv/vq21/wwpSlFr9lIuLxI3q7FVBEDlj5Q8ZDd88vSRUuypNsqyIXtTfnUfag/pHvTKd5Ui9MyFLfIQnLR9pGYlIq4mJjPMqwQz/6iWm5IlkBFUc+sBJevkU0R4Ez6Mf/ekqO7I57SaOkfbha6OZUWuATqe9NcV9wCnVKLxi0Jkq/d+VaImJjoixH16WtJ99gwNVrScK7gF5AuXiGwMFGbxSpxEXn1sVI2pKbsQTydr++4w9MyJ7tGPcqFwiQlk1EbK0Rk/ZCGo3ZB1YuQFDgDCqUKtJcKMxkiE6FpGJRr6S+JsKfXJwyrUKBOgp8IvV1lOQLTj7V9NJB9Vz8j8D+O98o9EMuDNhlFTbTnA/Z5Kfs3fpHf3zR/4v1YEbWkD0AjbuUHQFyc4pPzkZBMduniFA1CowK0J2fHIo5cAZlQCKN37rQj396lkEYecmh0EtEBb1aaPbFi3XgEzmspSRfcKHxW7lByWEtwSTjvgaLTMspHh7Tx58f+98YWcDIhCyLbQouIYmMSQtLLeGOlXxgryRlSTbFny3plhw4g8i6YrhGFn68zqKL6XVqhGgUks/ARTIP3N0AACAASURBVHtDuFei6GIOXpKoPjJULdzTlKwl+/2HZc/tbwbED7nzz6P8vlZPJmRC9gQ17lMmBJxlc6K7P61KYXMEXKYJfdSZ3GlIo6ewkiWVyAoa5BuMImCIlAtplOS2VdoROx1np2bkWe78pboeii5GpxLpWSVjXtVJNiiprk3ucn1w2xKLXbU5fKb5UtmXn7v+yoQs92eH5fcBAnSsm5KZ5zBdofV0lF2I7mSl6s9LyhWlK6RoVs7yA1eL0eNSQpYP0PTOkHR3T3jTsbuh2LG79Qzky0v35dROqoUMt+gliGQsLfS2XqcSEdXoCJuL/xDYXW+JjaGW2Q/ZyjHZJ/Xdfxvpv0WWYSbWkMsAHnd1HwFK31eaRmkekQJTROg1kvWBDdOrhXEQWYk7K9SWjkfJ31qKJTpcK158StMozXITIecaTFcOUix0P0xZqZy9JJHs1JaOt7n4D4Fdty4Tk/kgw6LNIoqP/8QZeURVY0L237NY7mci7TghLQd5+a6FYpSqD6zZFciZRmm94aRNSzHwiTlwBhnYuZLMiU4G4qL1uJwoPY2/QogadI2Q4GIgE9L4Y8Jts0GV+z9SHwOwo+5yqztk64hd1oqxQkTvK8ndWFhhu1nf44/hPl6Zd4ZnQvYOjjyKEwToD4i0Flc0SvNQUvWBtQ6c4erG67QqhOlUSHBBo3Z1TG+0sw6c4ep4Ugx8Yraevp6SCzqFcbWQzzUdc7MblKuIla3d9ptXFA1g8W2yUpkdhfAq3s48ghv9nzw7rGyC+6k3E7KfgC7v0xAhX0lyTQuzxkpqPrDFA2e4s69SC3xCxnN6rdKj6GhSC3xCR+n0+0zH1e4Uk8YfwoTsDmhlaLu19kooYJUP2ZwX2cf/PvX30DJI7b+uTMj+w7rczkRuThlZ+Ugrxeq1NHCk5APrKHCGqxsrpcAnZjcnCszijkZpXisdDYfqVG6deLiKk7vtzL7gZJtAL37uFnrpI/9rdoNyFzn3279702pT0A+rS15/fH763yHuCxuAHkzIAQC9vE1JhHylDMkWyCI2RKNCUrpzAypfYlta4AxX55VK4JOSAme4ug5qJ5XAJyRHVm5BmQzN6KWPCdmd3fes7ds3rimy6DKTsh/+7XNusGcC+7kXE7KfAS9v05EhV3JGLrJzXb/Xc4RRpQgdUrPyApZn2FngDHf2NdCBT5wFznB1LVIIfEIvSeFOfMFdWQ9p/HTszW5QrqDleZvNN6z1vHMZevb9b1AZevuvKxOy/7AulzORO403/FYD7QPrSuAMVzeYrIHVKoXb952uju+sHQXOyM41eCVYSSADnxT5gpObU9nDk5IvtjlvsjMMud4zBDZVXxeQSF3PXhjomcB+7sWE7GfAy9N0pB1TLl2KVOWNEigfWLr7JZ9VuqMsLXCGO2sMVOATc+AMV3zBXV1PoAKfkH83Eai3EnjQSxKdxLCW7OrOu99uY7UNfsqAbGuu3f8SE7L7u8U9ggYBMq7JyjUgOd09q9fSAAiUDywFzqCXijQvBvYIVOATdwJnuPowBkLj98QX3JX1RIdphC8zu0G5gpb7bdbHbQiIhjzgcn/3hQ1AD9aQAwB6eZjSWbxqTzEI02ugVBhLTEzh6bgl9dNpVIgKowQR7rtsOZPF34FP3A2c4Ux+63oi+sS0PI8stt2Zx9zWE19wV+fhONeuIuV+u7VVNllpyNaOxOaxSrLwKlv9oKtMyO7vFvcICgTIqjo9Mx/pPgqv6E8fWCKatMx8oe17u/gz8ImngTNcXTO9uITpKUqW7y3hy+IL7sp6wkLUCAtlNyhXsHK3zerKGwt9noryHZt8oHz7eUh8P3dFDUh71pADAntwT0r3rFeTsn22SH/5wFLgDPK1dZQf2FuL81fgE08DZ7izTn8FPimLL7ir66FgIXyX7CparrdbWZk0ZAv9WnX0bULk568xIbu+S9wyaBAgQy7yF3YluH9ZFu1rH1hvujk5W6evA5+UNXCGM/nN9f4IfOINX3BX1kPGb3SlwKTsClqut1ke85ZQhSl+iyU4iB8+D7/+rOtCBrAla8gBBD8Yp87NM+C6H44tfe0D643AGa7ur68Dn3gjcIara/Fl4BN/viTReitFaKHVqFxdOrdzAYGlMZuL7pD9yMojE5iQXdgebhJMCNDd8TUvugY5w8ZXPrDeCpzhTH7rel8FPvFW4Ax31uKrwCfe9AV3ZT0iT3SUjiN4uQKWi23ejN5cLIsTZXWi7E02aZCtPnunflRiXxclDGwz1pADi3/QzE5knJVj8HuwC1/4wFLgDAo04Up+YG9toC8Cn3g7cIara6W7d61KgWQ3Ez2UNr4gx0idCDLjLV9wV9ZDd+90gsFuUK6g5bzNoqgtMFIiicIUiuYelpSK5sQTXq5/MZkJ2fnucIugQcBXbk7OAPK2D6w5cAb98HuQp8CZuKXWezvwibcDZ7izOG8HPiFfcIr65s+XJPN62Q3KnZ0vve2CyC2l5DkuKT9y2b8fm/yM9xbhw5FYQ/YhuOVlaDLkSs3MQ2aO912DXMHQmz6wvgic4coaqI03A58ESqM0r9WbgU986Qvuyt6Qpb2wKVBaR39ypSe3KY7A6xFvl+CH7NsMEy+nMiHz01hOEDAYCnA1OSdgq/WWD6wvA2e4Co63Ap+QRmkwmF6UAlW8FfjEl77grmJTJSoEKhUTsqt4ldTutfAthR7HRXfGRcfWpAmbg4V4t35CWp+yiu6X/qwh+wXm4J2EtOPEtFzk5pctm1NZESqrD6yvA2e4s76yBj6hwBlEyJ7mB3ZH1tLaUjrD2GhdmVJv0n20Xqv0S8CR0tZiDkbCWnLZno65Ye8EREOemM6EXLad496SR4DeZsnfmAg50KWsPrDmwBm1ot40XViRs6TCiAp6Db74qyi5+a01FkNheYs3IiIqFF//Ysq1evTjsxj37IeFfpamI7jY6mH46BtTUIJ1q3/CGzOPQ6EwJ9tQoG69WOw58pSof33qV9i2/rTopyDjFqMRjZpWxfoPuov6qYOO4/j+P4VcohgVaN2pLmaubS0+vtz1Y/xy6gIURpPsNEqXQQ0xeGYjUT+r1SGc/z1VdKdaErTXrLvQZsgton5Bo8NIvZhZpKUojOjzTgvc+nCcSf5bDiMjJcs0d6ET6agrXS1bf/Cu3cjNIo3cfPyowjP/9RJ4pmXl49s261BAaTiF/Aphvdz4s6K0eP89uQB5BkNRb4USTY9MEWScbyhA1shxyLeKtqhWAPqlCyzzq14bYFpb4f4oVArkj19vqY/e3AVkfGiSD4BGg5Teuyz1sQcfgtFgXpoRRnUFXOuw11Ifl/U1FDG3QqE34cHFfQRmh75j64dcOESJHlBeqp+c8bT7wgagB2vIAQA9WKakH7f4pGwUeCeZU5lh8dQH1hw4o3rYGzDaKfpGTH/9IXR5uh7q1VjsQEYjfr8wWnzfqPqKwrd/q2ZG4PSl4eKL+jWWmKnA0oCg++3CC+LzPaJ/8WLEdxdN/R+outYcZbCovxH47LKJ1DrFFYUltJ5g/xXTC8GgKjvs5zcCa+N7ivoplXfbWRMTuc+4ZiLdxVU+cFAPjLzaBec2/4Hv55wWLxLWRalQYfDVp/H1tE9xZdePFi40t1GpVbjnk4G4Mmkbss6cs1u9poIWNd4ag+wZryH/+nW7em1MDLTTX4HqzeehyLK/NlFXuxnZz05E5NuPQ1GQb9NfUHO91khuOgZVDjxcSOS2e5dx+zBk1O6OuI9aoyAkFqr266EIjXWwT/yVKwjMDH03IBry1EwmZFf2h9vIFAEi48zsfL8leXAVJk98YM2BM+pWXAxHdjshFVT46o+huKW64/qaNSNw+Iv+uKvaMofuMW3b3YQFGx9F/epL7AiV1tV3aFNMmNwM91RbYUdYpDnsPfksatxYoVRCvnA+A0Ob7rTrT9rp/iv9sX32jzi69Iwjvsea+J5Y+cinuPx9ov3rgBGYda0b1tQ/gOwkW0ITjQuAF+K7YM8d78HoICcxveD0Pt8Lx5utgaLA3uiPNNomnw/Gv0+8Xqi127+Q1No1AZkjxqKgMOaxbQsjwpa+AfXcAfZrp4ZGIH/iBkS99RigUDpYnxEpz+5F7P62RZqzVasCpRrXOx5C5UOtxd4qGw4D4u6CMuZOVx9JbmeFwLSQdwNiZT09q7cs9oE1ZFlsk/SEDJSbkzMk3PWBtQ6c0bDacibkYgDTC0FZCbnbnz1xqvV6KAzyJ2SCR9XjGBSaUGePItc7QGBqyFZxEmHxOy70N/b151k5TMj8QAYpAmTIlZKR55MMSN6AzFUf2OKBM+6+Yblsj6zpDr1D7AZLJh0LjkYg0EfWnX9/HBcWf41LO39w+8haFarFjZslcmRNd/OkR9fqCGWNB6C4wXR3z8V1BCZptwZEQ56d08t1IQPYkjXkAIIv16kpQAMFzpBqcdUH1lHgjLuqL5eUUdc9zapjx+Ee4gWoNKMusqoe3O4gfv7yP0kZdXU984TlMfmu7ToYctwz6qq58yVL/0AbddGxNbEJHbOrOmyComI9qf4JSFauVzTbAnKH/GoeE7JkHwoWzHMESDtOSM1BnkEillwlLMWZD2ygA2e4swPOAp8EOnCGq2txJfCJFHzBXVmPRq1ATLiOg4W4ApZVm/FqImRToefBOhqeLz/Pyzd5Mki9sIYs9R3yUL609EzkGwyIjgx3OgK1u5aQgopR4dBpNSW2T/03ASk/xUN1302WNo5cXay1og219yA/12AKFm9ydsGgi0WuMjPi9go5hfWNsNFVYGahVS9NMrjqDtBLALn4CNVEocK6K0VaV98qW03GQIXeLEoo8NbVXsKlhnxgG+vWwIgC070VFFBpldh7vq/w06XAGbdEL4HCWCBchGh8fUgITpwdaFlf02orTS5EheOHRejwyW+menJzelG4OVG9yb64Ulwojn4/QNS/vfonLJ75uWkso8kNqfYtUdhx3HSftWTqV9ix9qdCrcvkiNSgUVWsPNhZ1C986WscePsnGK1+qRrcWxWLP3hU1M9/5gR+PPyfzX7d26U2Rq5tJr6b3OoQrvyeKrAxza5Al/F347GxN4t6R25OPVfchwZPVBH1Gxt8jIz4DLE3QgYAXbc0ww2FblAH7t6LnPQsITftDzXp/ofJYpvKyQffgSEr23KMTnvy8OnBUCooYEk+/uq2GIb8wntlBaCCAvefeEVkCyM3p+QhU1BgMBTGPgaURiWi1s0pWu+0YTAUFLlJKRVKKGautNRXWFkoi0hcYIRCqUbakK2W+kq7HrE8d/R8qDQ6xHe1cnM62QIFpNAXtlJp9Ljc5Kilf9W8Y0BUEyhCSra6dufvkAbOzMpGXp4BkREVbPY1WD68pNoeEA15voEJOVieIVmtg/6gx89ejWNfkD8rcOdtdbB09guoVDHS4TrWvrMfb64t8sVs37oJpr3Yz+EPQn5GLg4/uBaGTAMafzoIf79zDj/OPOnQ1YXuDU+88gN+2/y33b2mQqnEoPNdsLbrpzj/pb1lryZUian/dsGUpodw5VyqndwxNcLw2ncdMfTm3chIsz86r3tPFUw92Abda76DnFz7SFUdn7oZszY9hBvCFgPFLH+JeEdPux99htyJJtVWQ6Gw9YMi8vnm4vNCpobVlznA1IgfL44U3zeutsKBm5IR3xa6Qd1fdY0Dq2sjTlwy+T23qbrOgdW2EUcvPSfq+1QhF5Ji0aOMRrwdbyL856rQj599/br4J0W9IzcnIlbzC9HSuPft9pbqR1ztIvq/V3e7vRuUAuh2pieu7vgDfy48YVevUijQ5bfn8euU/Ug4/pMdfiq1BnV2j0TWa6uRddaBG5Rej/AlU6BYPAWG69fs+isrVwFemIHwtX1hzM+2r69yB1K7T0XM7g5QFPfXMxqRV+dRpNz9AqqcbGF6CbQpRqTf8goyYh5D3I9NYFRFQNniEBSaCLt53P07vHotCbPf3IyT3/8mxqp3842Y+MLTqF+3pqx+f5wJO1ZJz6Tl/dnS3Mo73Cf1bxSYnnmpF9aQpb5Dbsq37t0D2LnvOLYsnQR9iBbPT1iEWjdWxayXTVpb8bJz/3HcUC0WDW+7Gf9disfAF+dhYK9H0e9J0h6KSvzxf3H99AX8uf5boe01+WwoDtz9HvIy7C1nydWFLGvX3bQXBWYNyGosIr0hF7thSuX3Hbuq/L8mOyu+K56L3eHQTYj6r4vviWdit5bQH9hytRc6VtnkyNMF5PxyNHEgboshNyT7cIgarQpf/jMEjautdGh1XbGSFod+HIg7qy112P+2BpWx9dCTaFxtucP6p5+/G6MnN4MjQqa1bTv1tHBzckTIVH/ssklD7xP7rgMjLhMhH9l4HtvGn3ToRrU+vife7v81/jhw0e55MFtVb2n6EVLOZ9rX//+Vxcj4bjjY/ANkJzgICFMAdP+rJ062eguGXMexzdt+Nwj/6/C6OFUoXhQwou7esUh4bqLDvSc/q4prX4NxyhDTqUaxQv0Vs1YjbFVP069+sUIcnDlsB0zasYOiVOJ694Oo8lULh25QUOhwtdkxVDndxHSycdMgILoBlBUftBnM3b/Dl2etQnJqOpbPHQ2FUoEZb7yFawlJWDVvrJu/ANJuPkYQsol+zac2RfRsFfXFy/WLClhDlvaTEaTSPTFoGkjLHfR0J7HCw8e/xovTV+DnTza6lEJuyusbcPHyNWxYNN4Goaz4dHzUZp3puyAg5H2X+5VImFIh5LZV1zvwKwbeOvWUIOzSCPnQ0j+wa5a9VTMRriuEvOORTxH/Q5JDRhtRRkK+76sB+LPLG0FByASQsuVxKLQxNli5+3fYZ8Qc1KxRBXMmmE4/3v/wBJZu2I1jOxcF1S/VC8ptVlbWpmh0lrQSFhcoc3Yn79UvZkIOqudINotp0mEoZo8fKEiZyq9//Iseg6fjy33LERle+r1UXr4B7XuNw6Ntm2Ps0KK7wCuf/oNzu3/GlaNnxZh0DyyHI+tuN76D3DzHR9bDF92PJjeslvSRdduq6+w14P+P2EUaskqpxFOVKHOOv4+slRhx1XTH7emRdbMv+yNh0REkfCLzI2uzhl+5HZTVu0AR84Dld8Ldv8Njn3+PkZOXoG3Lu9GtQ0vMX7ENA57qiCc6tZLNb48rgo5QFlpZC7sIcw9xye/Tz8sK2Mralf3hNl5EgN4273iwP1a8OgatmjcUI5/99yI695uEI9vfQNUqtm/xxaeetmAjDh49hQNbXkNspShL9bVT/+HLAbuEYQ/9/Df5tCi2s1SNuszCd6q62aFRl7m+cfWVsjLqurN5VWw63BUqlRITux3F6Q/P22yj1I26mn/Z3yLvX90Xw2CO7lVo1HXz3jGWeqkbdVU+3UTwiDBWvHsTlNGNCw+Q3P87vHjlOgaNm49bat+AL775GSE6DTYumoCba1X34i9E4IcapqRrpkLzart/zafZ3q9fwUZdgd/88igBvZnTsVe7VqYfB1c15BWbPsDyTR9g26ppaFCvlgU6InkKkVnc/uVw8+2IvD0GzdY9ZAfzP3su44spJxHXLBYPrbnPrv6Hzf/gwNRfcNuj1dFtuSnxgXX5dPnvODTvDJo/XQedX73Drn7nvDP4eOWvePj529Bj/K129Ztf+RZH3/kX3cbfga7DTb6iep1KaJW0lhnDPsHnh89j1NTG6Nj3Nrv+L/T9ED98dRFTFrTGw11MFsnWZXDXA/j113h8/kcRuVjXD2m3Hxf+SsOBvx2/lfdttQeJV1Kw/0xfh49o/3s/QGZ6Frb/4rj/i3cfECyw8DuTtXXxMv72fdCGqTDrVEeH9bNuPYiwOB3GfErhIu3LktoHUOnmCuj9kePAF1tv34PwOyqi0/aWdp2vHv4Lv776NSo2roKGC+zHj9//K66uP4aoe+vihgntER6qMSWOKMylnfHhV0jadxjhrRoj8knTtYtN+fh9ZJ34DPqWDwAPd7Ov/2Q98L/Pobm3B/LutV9/yOcLkPvfKSgaDYGhnv2zqz89HqrU/yH91heAOPt75tAzQ6DM+BPpdx+zmZtCrupD1MK6n4q7f4dPDpmBVi3uwrBnu4Ass6ct2IQTp37CV/uXQ61SOdwnOX45VLW1pCtjn36/ijVkOT4u8peZ7q4eebApnutt+rF2dodMLkVvrNqOHfuO463FE3DbLUUuTWlnE5F8+hySv/0PN03tIMbbU2+j6TXWYgxlxC0DG6P+OFNs37XkpgSD1X21EYMvFRHL5LD3bEE2ArMzHrd8NzLcPgHCkrSi4/Pe4ZQtpuiYlu5Et6YVhcXrIeqLpqD6XWlPQ69TQ61S4HbtMjtToK+STRr/6oU/YOnsE0XHwOSqpFThl4Rhon76qE+wc8svlsFpbPoB/jnRlPzhxa4HcOr4paLx6X5MqcaJJJNBXZ97d+LPM4mWenrZ0WrVOHHNdG/Yu8FOXPgvxao/oNFrcOTKs6K+/007kZ6QU7R+oxH6iBBsvmhyAxsauwtZmXlWW2NETPUKeOOMySr6pajdJjcjMz5G4KYm0XjhExMxvRb9Hkx+PqYGZHRzU+s49PrARLxrq5mDOpggIPlv61sP98+7W3z+qOFakJOSeXvoRLLu6Ka4ud9dov77h+bbxKum/q2/moz8/AJk5xpwvm9REJDCCXDjlqJsTqnjKF619RG9ERHzN1j2I//1HjZ7S+Orxxd5EKjXtbfpT0PlDThk6R96mNzFCsenR1xpREa7k6Jed3ExlP9tsrI6p71VIKPp96JerR0NpeJ+aJVNoVDUgTt/hxmZ2WjacajwhmhzvwnLX878i55DpuODjbNRt1YNi4xy/59BSoplXXg37Md/1xj4yFruz44s5Sc3pl37PxVW1qF6HYaOX2hjZU0GXtXiYjBuqMkNYPK89cKAhKw5a9esallzlcrRUCmUONLsdcG/tx96AWff/AXnNpt+gIqXNt8/g1Ojvsff+/+2rzQq0Ovvx7Gp7We4+pMDQyGFEuPjH8OsevuQnWh/56tUArOudMfwGu8j38GdsC5EjSXnumNglR0OExRExoTg3ctP4k7tcodW13Vvq4jlR7uhWfVVDux2gQcfuwVzV7VBM+GTbHtnSz/64+Y8iB796+GRuA12hnNUf+iKiZAfpjvhYjMQaR0ptJp+tOpGOzcjIsUDl02aeP8q2xzUAxuvmvZyROWd9m5WMGJZYTanaZV328sPI2bGm9I7LiM3p+JW50Yjhl/phn82nsfJ2fYubmTg99TfPfH71GNI+Phfu70n1/EWJwfg3ItvI+fPeLt6JZRocWI8/n5uOvLS7F3cSPOMXPUqMHMkjPn2Vt0KtRaYuhShK3oBjoy6QyOQOXA1Km7rDCO9bBQrxop1kPzIYsR91tKRUTZyq7ZBUt0ZqPxjU2GdX7yk3zgeGdGPIzT8NiigRYTmf1AotHD377B9r5dQ68Y4zJs8FKEhOuGK+MmXp7H3rblBpSEPVL1reZ+3NoI3G3YJfItfKVsl9PS0fr2BY1nLktDkLjS9bY+buRKfnfxRLOWOW2th6ZxRljvhbgMmC4JeON2k1dEPwYXL9v6c+ye9hJwfL+O/7d+Jdo0/GYNDjTbDQPlsi/+oGY1oe7ov3r11NxT59vX0o93r7x6YV5mCLjjwRTEC4693waRKu+HwV+//U9fOie+OQbE7HLohkaZKGYv6Vd4BpdJ+/AIosD2pF+6PXi20muKlwKjAkSsDcK9wc3LkBqXBZ/8MBAUJcVQfVSkUh37siw5VNpDKbDf+zQ1jsfTQo3gozpFfMdDykZqYuvEhPBq3yd7Vxwg8PeketOtSG+Ma73NAmMD8bzvh96MJ2PbyKYduTsuu9RBuTn86dHMyYua17iA3p9T/HLk5ASOudsWuhu8hP93B3hcAvf7pgS/u22TSru0fDtx3cqDJqrrAsZvSHYdextX+42G05Im2GsQIRK95Vbg5OXxbQgEUM9dCv6SXYzcphRGZI7Yh+p1HHRrIGY0FSH76IKp8Zn/8TlIYoMT1Bz5FlR/JSNJefqj1SLvnCyg0lFNaAa1yALSqlsjOutetv8Pf/jyHlZv34OiJ7xGqD0HjhreK4+sG9WvL/SfJRv7+ysITLnMCZD/9u5EJOaieI9ktJiUtA3l5+SUGBHG2oIzzifjy8bWmZgVA229exp77tiLjUrp91wKgzQ/PYFfDvchLd+CbWkjI86uSX7IjQlZg/PXOmCT8kh05jwJzrnfHkNgdDn8TiQbWXu2J/sJv2b4/5fQlV5+HhQZrX68KUePDf54tkZAjCwmXCNuR33K9BlWw6VA3hxoygfXk8ProN7mZQ0Km36NNX5vcmDrFbXLo5mRODtE/llxG7DX0jfEmH8uRle3TL9L4RMgJ5zOw+J6PHBL+jGvdcHL27/hu+e8OCZU05COdP8b1X5Id8a142fr60a3IS8xy8GwYcd+pgfiz+2KH6RkVRiPq7h+LlKGTUGCfjFp4o0atfg2YSn7HDgppU7NWQ7/sSSiM9jpsgcKI7BHbEPVuJ8d8DgWSeu9D7Gf3Ce8Bu6IMxdX7D6PKDya/Y7uXOd0NUDY/gLT82pbTizDNp1ApbhBN3f07pBfq/PzgjdT1rOodZz89Pql/y8D5kH0CLA/qewQSv/kX/27+Ggkn/zFNpgCaHX9R/CDtqbfeXlNQKdDmmz7IOJ+Bva0plKStjLrwEHT/oRP+PZGMbd2O25AKEUZ0rQgM+eZBHF/6Bz6e+T/bHz6jEXWaxWLAvpZ4a9TX+GLrvza/i9T/vl434dnFTTHnsaM4e+qaXf+npt+DdsPqol+T7bj0X7rtHfP/G3zN29ERd7esirb1NiKzeOQvoxE7T/YVhHlfjdUwWAffLbxnPVUYuat93AY7Ddr6yPqhuPUOIncBR66YAn10ittof6RcYMSBq/2hVSvRO3qb3QlAQYECmxOeFOFFHRJygQLLrpvumB0d+39phQAAIABJREFUWZO13ozrpiPr5XH0QmTvRkWETGVbbXohsq1XqoGeZ3qIvf+h5za7enWYBvd+/AwyvruMC9O32t7xAgitUhE3rOsH48HPkLTnoO2RPN2x170RFV4aBuzbgIJvbI/M6Thf2aQZ8NgA6HZOheryH3Z7r7i/FzLu7oKK+wbDmHbJ7o8n+6G5yIq9E1X/r72rAI/q6KJnLdm4G0ELFHdapFDc3d2hUKDFKS1SKF68eHGCOwQo7g4FijsECMSNeLL7/t73VrL7XlqgIcnmn/n6fSWZJzPnvsydO/eeey80gjbVdLNJzw+p4AfYFYD7zSpQyEzPxOlTSPrqBj+nZFlhw3etlNWEtaI3lPJqn/6P1cLe0FWxgZcxnxRE70NOkwREnywko/v9mEK2sC+FDdeAQOjZJ7g5nIJhZPwRbImDQipIb1c1Lv1yAc9X3xOsGeq2tkKti8YsONfHXMGDbQEGsr+1rRXa3BGCiqjt7HYFj/94q9PzHNTuagx5YIxm/b3xOby4FqL7EwWcfW0x+oYQUEZtYvUjePswUqBgcRx8irhg4tn6hv5R5fwRHhhnoFYUq+SFUf7GLEqtS2xCbDhZckLu6Mr18mOynzHatm7hlYilwCidOdakXUmMW2A8zqyWfxVSUoR+GkOP/iUxcIIxkrxx7rXQ8oU3OFBe7T4Tv0Kb/nScKbR6uVbxwVD6vN0/raiBWk0LGvqb+qzVWYK0YCmw/40Qie3uZI34JA06Om8y+EIpBenmyE5QqxSI1J1MDPbYYXIKsDCkneHZN/cHYXcvOtbmQ69AaSx/DjbmFb+5/BUuTLxOmb95/GUKOQa+Nsruxqy7eLjsPk8jIyWkslahzT1jpPPTaWfxdp9RKarUVqh0sqvh/SFTdyHyorChIgRU9rb4bPMAQ3/yrN8R9/iZLmOlDNauLrCbPtr4l7nkF3BvXxt+lvnkBgZOMPzssH4QUqOFVKy8dHMVQ2Kbnw39Trs7Qhb/zpARU5unIt5Vn2To97nQEJqUOB4/YhWkeDdBVJEfDf2et6oC2hThu4ccib49oP58OEKikvgNkb19SYFiJwPslWuglLPyjEbhCf/qrCDuvDE/uL7fxIf8Cfo3aYzfofmYstPPLHVmdpJGNhiL9u8gntCoJGjMeE7XG2+BQ1EXfPF7U0TFmgZehe8Pws3pV+BZyROl5gqFDdI2ynl9ZdYt5K2RD1UXCNHYaRspgrNz7qFUxzyo+YuYhnRkzhNcXnEflfoVQ31dYYS09++ZcA+XtjxH/RGFUVun/JQKOdydrPjFcul3l3HzcCA6jy+P6l2E3MCOtirQXGMTUjG73yXcPfcG/WdVReWm3qLx/dLhJJ7cDsf6e8aiFmkvIss8+FkMfrsrQcMBMLLeAUQEJ2D1Len7qQhEQqwWc/6UpinNr3ESyfFajL4qTVOa89UJXiGMuijdv6z0Edh5WaHb0XRoTGUPw7mgLRrtlPajHvl6P5xy26DSJunnX2/hB7sCHigyX0wTirvwFEGLTsCulBe8f2zO4y6XyxCl20AknrqOpL2HYFeyCJR92sDH1QZBkQnGKkDnTgGnDgJlqwBNxfjKL/lDfXM/Ess2hbZyM5HsHK8tg9Xrk0go3QFxnwknAWmb88PpsAk/j5h8/RHnK77f9dWPsI66iZjPhiLOroHJvW6O1lCeHgp57AO8rb7XpI+vJuZsLRlzkA3+zLNsCB0VflkSZb05lSnkLBM6e/HHIRB98xXe3X+Ft2efwHeyYFldrLqGt4bTMGVQoG0p5BrxJd//R/EtZr49DqWGVUCefoJVuDbfLtNALg7o+dJIc5rtvld0zDgyzGiV/ehJNCnTY9Lpuqhgev5QTyOthX4my2tBSFvQYklUmlZOfiLfIeW5ppbLzQZVlVTgwUh1IdtwR7Dgb9ow4y62z79mOEYly5Yso93B3fj+hf0u4cq+52n6qRiVDOvfClHPPzQ4jId/BaehOQFKpRw7AoXFYdQXBxAa8C4NDYmDtbUKy18JintK6QOIepuYlmEGOnGYEiAojomF9iMuOslIM+IAWzcrTHog9M/NtR+pyakmNCn3Ek7ofVo4MVidfy+0mlSTY2LfGl5ouEGw+PcX3wrerauHhwN8GufDF7pN19XqK3TWtPC9ET65u1aA7wCBA/+oyXxwfHEOo/xqnP8RIdGCRRnZdxx/v6FxHPJs+BVWKlLYKUgZQ4GHphw21czFhsvV87uYHYFzSBxKBTeE5rK5icnc6GAisvN+Q7/nia9N+ukzCKp1mu+3e7sG9kHLjGOjEwO5DEFlLvO/c3swF4oAf5PhUcBfUD2Bn+x2ZwhUvjUB58KQuwtZ81gD2iv8sgSGbRrhbza7N2YhZ3cJZeL4tEmpOF9vGr9EFto9CkFbHuDZAiofaBZIBOCri31wc/QFvNlvmimKH65MjsZ322Nf46OIuCumslDwTI+AVlhS/AjiQ8UVeWQyOUaENMX4fPuQmpAqQkCulGPqm5YY5r1Lkspi7ajGmrA2aKRex1vB5s2noBN+vdAY7bw2SAaJla/tgx831UEL73XiakcA2g6tgK5jSqAbH2Rl9nSOw7CNtVCunjdaeq2TDMLaEyzwinvwNCbx/WtDhA3DKH6zYe7TBWaFChua0R47JGlYs0IFhT7TfY9kVPZo3YZnVR7xZofQ6vuqNUKOBuHyd6ckfMpAswcdELjsGgI3SlDgOA5fnuuP1xP3IO6aLgbBbIpF9g9D1Lh50ASFiWUrAz7fNhtBI76DNl7i25AroZw2H+pFPSFLFVPkZCprxA9aDZetTWFWqIt/l9zGBWEt/eB98mvJXNqcfWGEfLkS3jcrmdTq1Q802a0hovJNgsehWpIBeHGlBiMuVxt4nKnGuyzkLS7xXHTWBATaKNZniYW8I5UpZPYNWhACERceIezsQwQfuMGbmQX3jOIjZ1MjxFQY6q97oz92F9sKjUR5Q+pvfL8j1uXfKbmoUT9ZybPc90lTVThgZFhzjHHfJckbJljJSh7iuUMycpasumXR7dHOeYMkzYnuJyu5rZdwfGbeyIe47W03tPBcL/l+lY0C2593kVbIZF25qrH0fku08lovqfALlHXD3ENN0MNzs+T7G31XDB3GlcFIj52SQWA//ikcDc+ocEiynxT26SkPcHnBQ9HcyIr9Iawl9jQ/jfAb4eIvlOPQ53UbHK2+B4lh4tKWJNBmDzriWu1V0OrTXqZ5Cu1/Kp37Bo+bzje1fg3XcCiyfzhfzUkKHNoQ5Fo7HXEjhWQsUk01YxHU8ztLbpZofAnDNomsY/1z9Fay1/GvpaOmOSC0zml43agk+W6N3AqJlS/Cdi8FbEnRoJwQUnsPPM9U53tl+VsBuapD7mXMc21By0KGD7W1cr1QGz297Jif6Pe7NdJZ8TJ8gv/xgcxC/o8A5pTbY58E4Ubv3w3TKbh7FO723I3oh+J6xfQH1fbRIKwvsQUp4WIrhvqb3O+I9YX3QCtVgk8L9HzVBrM8/fnALLHW0Cvk3Xy2JPNGCndGWGsMcd8hrbA5YH5oW15hStGoZHJg/dtOaOO1QZLXbG2rwoZnHdDCi3bz4uF553XAsiut0lXI5RrnxrA11dDSe72EBQz0nvElmvUqKqmQSWHOvdYcbnntMMpjpySveLbOQqZ+iaBo3oImmtPK8sckaU5kIV+a8hfuLheKhaRtJI6+r1vjTLujiL4jlj1FOzV72BG3229DwlsxDYqC3b4ghdxuEbQJYguWNmNFDgxDWL8JkHFSmTw4uK2chpQfvpOmwFEwGCnkeV2kN3NaDgnDSSGnR3PiENHpADyPfy25GdIoVQj7+hg8r1OeajENilM4Q17jNDRbK0tuCFKdSiCi0iJ4na5mUPjy2lsgcxKnYM0pa8eHzKOFQjg14qs8ZeL/96QyhfwhcmLXZiECMbdf4vmy44i5/UoYhYxDwV1CdOv5yuJkFkT3bHh9AOJexMK/FhWxN20qOzXqXW2JkLNRONDluHm2QzjmdkDrC/VxeeEjnJn0QERjylfFC+39K2HPkFu4tOmJqL9y50JouaA0FjQ7i2eXgkwXVo5D40lleZoT+WiDAmJMldrfRejH7qyFotV90L3QdiTEmluBHOZdac3TnNrk3gCNKNEJhz26VJbdvSkvrzkvGPDT8YJbea4XJTohhbcnWFgcJC1kLbA2tCNPc/remTYcphsSuv/fFPLc8LZ8UN6vvH9erHApCQu1lbl3SSg1Gfq8EoKn9hWV4D1Dhub3hVSml6v/LpI9J1Og0pk+PA3q1bfknzd9v9LOGgW3DgTuPUHYXMpMZmwUGW3t6QKHaaMgP34YiUfoBMV4BSFhlb8QMGAobI4sAXfvnKmPmPy8xashof5AOJ34AYqgu2YuaA7J5bojrlgHeF3qAMQHmW5GwOFdpVlIsKsEn9t1oU19J+pPLrcLWtsCsNr7NeTmvGmOQ0iDU/w9nmeqGcfmXgHyYv0hcxfSYv4/t2bKdVkyff9UwU2U3RuzkLO7hD7x+GinGn7yHu5NFI5/FQoF8u8wFkV/s/A6Xm68wecoptVNpVLgy7M9+VF5uahxov9ZBOx7YigWr7RTocFVY9DW6T5X8OzoK4EGJeNg42KNjjeN0aybGp9D4JVw3eLFwT6PPQbcMEbzzq9+BMEPjQujVxEHDE1Dc5pQ7g/EEM1JF9D1WSU3DElDc/q26G68o4QVOq5NmQa5McrPGE3cJf9WJCWm6NL1cajeoQi+XyAErFFrn38DkhK1OkNbjuaDiqP3BOPC2ttnG1I0+vzQMnT6pTwap6E5tcrlB2iE0CXynf+wiiK5jTSnHj5b+JSPgt6RYa0uDSbRnOISUvGd83aeSsP3cnL8GmoaKfyDx25wOn6sjFNgQXRb2ForEPEuGY95mtNVUDYq/n6FHKODjdg/WP4I56dQbm7icAmy7xlgDKgLWP8Id6bdhIZoTkQzU8jQ9I4xrzh9G2+2/gktjU8mg8JagQrHBE41tbC5xxB24rahID3l5S64fbChP3HhOsTefMRTrGjuSkc7OM2jo2yhcUvnIOXFMwNNysrFHfjBSFOyWzca2ojX/LdHz5C75kZcj18N97vt7QVNHFHohNNlrVsxRNc35sb2udAUmoR3/PMp1kDjVgURZWca7ve5VROa1Hihn/Lj+DSFstg0PnKf/m48j9WHVpMsYAMgsUA7xBQ2HrV7nqsJaCgGQgZF5ZmAT+33qkn+if/ks/TxjRWUjS7NJktnKesrS+gtZ8M3kEH9B5hCzlK5s5e/JwLa5BSExJpWc7rdai200Qm88edRvwh8f6ht8rSLdTZC8y4ZcpkcBbqVg89gUyqTf+ldSE3UQCsHSo4ogyK6iGv9Q9Z+tgfaFC1pANRfVAm5zahGv+XaD20Kx98/PI0C0d8/1ZuigznIFTKMDRIUCFmULg5WOOIXgI19LvDRwTKVHL++MSoYuu7E8kfYM+EO36+0kmNeoMDDdbBV8tbu2p9v4sCCe3xAkFqtwOIAUwW4adRfuOj3GEQ1trNXYdZT0+f/3u8S/vQP5JWgvYsac3URz/qxL+1wHs9OhPA8WWcfa4y7ZVqxaWmzs3hxScj57JzXET/+Wc8E+1W1jiP4jrAB8SjqiH5nTalM6788gfDnsbyC8C7vjo6Hq5rcv73cIbwjNwNFRNf3Rf1Vxs0Hj0+V3UiOSub5574N8opobJfrboAmIQVyjkOu1iWRd6Tp/YFdZ0OTROBr4damOnL1rsmPJSpOOL4O7DEDXHIqKD+5Q4e6cGhuSpPT/DSUly0dFys7dIF37WqGqGy632ZBN331eiQMEUfsOm9uATkVN1GoEN5+t8nc7XAY9i/H8Mpbq7BCiO8V036KrH65gj+V0MIGIV+YVnRyj/kVioitfLlOmdwGb/NdMLnfOWosrJP9wWlk0CrtEeZx3tBPGVW9nCne0eY9/zJz5mWNlALPXtjCCKc/mfHzoVTBiDBvYRHRsLO1gY3aKlsAzizkbCGGrBlE1KV7iHvyFm+3nkautUIChBu1F4t2sNb2Nii+TyiQcLYSVdcxpaI4FnJDmU2CYtr5+TaTydARq9cXPqi+SbBKV+bdKcrGVKx1YQM/ea6IBgV89XNRVPpOoFFN5iOH07yCA8aFteQTZ+ybcRc7J9wUHWHP1kUdbx1yBec2vTQ9Aif6UogQlTyr9UncPxcsomGt1h1BT/ryEIJfCBaVoXEcFuuKNwz9fA/iSZmZwMNhRYhAg/oxzx4k68oM6u+ntJ6zdVbvGO+donTQ5JOdqfMZT/HcI3KJ06ZpXIiwqZhHucLNUoPK5QoMCRLKGK7OQ8UlTLFTqeXo9li4/2AxXa1a3SX8EbGjGvUuCf3nq5hmGiNrxiGPK0pvEzYtz9tOF/lVHcvmhfuELjzN6XWnKaIjdKc6ZeHQRxhf8mhKQGN6xu3UpBnUTZrwiU9s5nYVuR8ShgupGO3ub4XqBlHcTL/NyC4Czck5chKsYnaJ+oPzCznf3e8PgiL6hvmnjZDKgtL1ftkCSAkwGR9tqoIL3hTuD6kGucb8iBsI9RGe75JYGSq7/pDbVYRcZbqJSSORHP/P+kqhgIpJRq5M+PmwmUJ+GRjMF94JeB3MY9668deYMLwHVErpUpczF2/G+u2HTeRTrmRhbFg0NkNlxhRyhsJpWQ9LjYnHjXZT+D2q76YJiD70AM9nmVoFNCNaeMudGMRX9Ak9TIuSWeM4VL/SG4frHUNsQDpBYI/bY0OpfUiKFtOYaIvc51Vr/OZ7UODNipocw0ObYpqXP7RacSAQKZ35SW3RQ7VFx3s1fYBKJcf0wFb4js9iJX66naMVpj9pjl6eWyX7fQs6YPKFxhjE59IWP6Biq4Lotaw8vkkn13a3uZX4hCQjeE61GLs5ug3DaIkgLrJkf9X1T3XfI0m1GRumV8jSPuPhocJmaTXRnMQ8LfR+2QZvdwfgxtiL4uGBQ5N7nfB8wgm8PSamMfGpSy/2QeCw35HySkL24FBgx48I/GY2uFhxAKCckyHX5rFI/XEEtBoJGhMpw2XLETe5A2RS1ZrkciQOWQ/nzc0lAwRlSgUi2u+B96syksUt5HJHvM1zBh6Xq0rXNbHLi5BSW+D5pJxkRH+SQ31Ees6Ex9sykv1xjuMQZ9cOLjGl+I2idd7rkMkdLGuhyMDR1lGtEU44DBay3lLW/V8ffp3B/cdTBYNC374ZNRv2djaYOqYfgkLC0b7/JEwY1h3N6pueKOmvn7FoE169CcHogcYyjpQzwNvDNQPRof20kMePtf8zBCIvPkD4qRuIPHWbPzPy3Twet1uugsY8lzOvkIFyJwbiTOW16URFc2h46xv45d0kWQCAj8p+3B6rcu8GpKKmKbL3VWvMcadqRtJR18PDmoutY53M6PlTottgsDMVl5AQ5N9pDWeHtcX3fPEF8QV0fL0wrC1681HZEg/gZFgd0h6DPMXFG+htdPS94FWbdBWyg6sV5txviRFSCvdvRZm3vDuGHKrBR1VLvb7294XwRfdCWFr+qGRU9bfX6+HO+le4/puY5kQCGRbagqc5RdyUojkBvV+15mlOUhHzFFXd+EEnXKqxDtoU8WaJlo+vLvZFQLvpkrxewif/jjEI7DTFPP5NLz3k3jQeKaO/l/x2KArfdt5iKH/tLMkyItknDt+gq+YkERUNDlGd98PrRWmeHy/aD1G+6nx/weuy9EKs/bs0aEKNS7C/U1Ka5qRQISj/FXi8LS1Nw5K7I8TrOFxjSvHfltyuOZR2DSG3NXVF/L8sPzWVdMqS+VHWJ1OMCpkKflRtNoi3bsnKpTZ1gR+CQiL4ynhSjRRyVEwsZvz0zScVFVPInxTe7Pvwd3cD8GD4csOi+F8t5A5PB2FNef8ss5B/Cm6GUXyAk7gEILOQLdtCVv36G2wXdJeuZ/yJLWSZXV6g+h5w14pliIVMf3BWPjsgty6TfReHTziyr1UrdT5j40v09rL+N5/i57MpfQ0vfPoiEM17jsWpnfPh4ebM/95vxxHsPXweO1YYgwbTwkAK+cjpq6hcvjhcnBxQu1p5VChtzFGfUZAxhZxRSFrQc+KfvsHzOTsR//Qtfxwtt1b9Jx+yfSEXVN/djs/buya3aXm1zPAh/xzRiqf5ULUo/19ufbQPOb1qUZnhQ7ZWyTHMbXuG+5AVcgWGhzZHqkabJT5kmxK+8JokZEn6GB+yVYMmQJ0GvDJUZ4EPOaXeNWg0HGweNIQshWiBaWlYH+ZDlqVQMJ4MMqtSULlNgty6lAWtGhkz1Koqosplvg/5fEo/wwRu3HmMroOn4oL/Yjg52PG/3+Z/CsvW78WJ7fMkJ+p/5AJevA6CtZUKdx4+5+tWz504EA1qZmw8AFPIGfOdWcxTSAFHnL6D59O3QE40pHKF4TTSWK2JJvIhUda52pdG/hHl+Pl7OqsR8S4Ju0vszJwoa6UcM6PbICbe6Hukikab+lz84ChrvQB3zribJVHWns7WiI5LwfyGpzM0yrrHieqwt1EgPEaoU53ZUdZ2HYyVsOj9HxpljXJCXmxq9jZKYGYnofCJTPbJo6yjqpziI/f1NCfHsJmwid32UVHWNH7XlIFA8llQIJ/SfTF/bC2VKc5iFpOPGGhlFSUfMvch6x+Unm/5v/dfSjEeNest5NO7Fhjqxf+bhWw+1THTfkdU9Dssm2mkiH4EHKJbmELOCBQt6BlcqgZB0ckm/rpn3ZZBGxXP5961rVqQr8qTtt1qthKauFS+36VxMeQdaZoG8EzVzSD6FLmHCw/6Ej69i5rcv6/4Dp6OKVNwKDe0DPKl4enShSsL7AVH+kKhRc0lVVDYjAa1wPsANBoNz5MdEiTQhPQ0J1osf/XcBzItFSoFhr8RInb1jRT0tj6XIdNqILOSY4qO5qTvJ6v64OQ7/EmB0laF6QHC3B1sVbzz3G/ULVxd9pD3o1s5KDDxmSk2+wZfw/Vtr3j/qY2zFcY8Nq3YtLHbFdw/FMhzZR291PjhjmlFJL9mZ/HqUiQvD5e8NhhoRnPa0uACgm6G8XaZSz57dL9iSkHbUfsIoh4J6U1dijugzRHTikx7qxxAfEgST1PyqOqN2rrCEfr5X6yzDcnRieBkcuSqkxufzzClUd1uuhap8Qk8DcyjkZgC97zrAiBOiCx3aVIR+b+vz5cu1G+SogdPBFKE782peUPImph+O9y4IdCmanjfuLx9T6RVwDRG1bzekFNyFk6OpNGrDZu+VL7E5d/UsNUdwHGpkMmUiOq91UT2diFnYH+NfNcUCKhASKMjpv1vNsE+YIlAg5JZIaTySZN+9+SlUASs4C06uUKNN58ZaUx0oVPcVKijt/PfDqd0QKgX5X03NtekEVAkHRWCyZROCLM/Y+jkaWmuCshk2YNuk1lLWEXVMmETosufaeAdf+Kfryb3N0xRyoc8ed56hIRFputDNsdn/ood+PPWI/gtNPLmMwJDppAzAkULeUb0tXuIu/MMEfsvwHX5eH7Uj5rPMfWN/R0gpXRUo8BGqrQD/FmLFqy0VBLAJrcLim0QrOqT5c0KKHCAa1lflF4tKIZdhSky2QgQLaMFWudH+RnCUc/y3KaZvugPtMqYkiitoznN9TCvBgX+CFafOGOio1m2KQ7QF08gZXv4l9tmVgiHaSFC4pKtQ27h6qbHZoFSHGbpaFDL6p9BwI1Qs0AvDtN19/9a6iDigqgak2kg2MQQIdPVxPwHkBxnrMZEv6OYtSk6mtP0XHt5vrVpk2NsqKD0F3jvBydKLynD0GChf1XePYakH/pnUI3kPgFC1PXmz4jmZOpTV6iUaP9A6D9ZwVx2HFROtvjqhFDp68+aS82ypHGwdndAyR3CEfTTVpSEw/QI16lQbnjO7sIrqchvxphR5ACHmhWh7CLgnzxmkFkgFAdVveb8ETU11czuZhQ8QDF+gyHxieOKNiKK2bt+VF0MsHuxFbb3lou+3ZAmx/l+5xfTYR24XxQoFlxVULpegd9AFnvNTDYcgorcEPpDm4FLCjD5NkinhObW0ahiq4FLiTaZPyUfiXSlRCmAa/IXkNv0hFJdEXL1/0/d5PJWS80wzZwfryd/a/KiviNnwdHeDlPH9BVFWb+LjUevYTPRp1NjNKot5DSf9/t2NK9fFXlze+Ph05foNXQm+nZugv7dxCU7/8uMmEL+L+hZ2L3JoVG431egOXmspv8Dj5vPkZgFh8L7RuL5lBOIPC4Vucuh/KmBuNJqL+Jf0qJj2ohjWOvPHvij4h4k6I5KTa+Qo/WjtlhdZD9SdQkj0vZTRqt+r1tgrvd+3vI1b6QApya2xWjr7ToKhekVSisZhgc2x08eZrxb3WVqRyUmPGmO0V67eMVh3rwLOWHEhboY40m8VXEr1To/Oi8rj4mepokn9Fc2nVceFbvkw1iJ0pH0vmk6XvG/0Zjme5nW2OWfzwFDQ4QgrRUSqS9pNv1eCbzgzQW3iyt1cUDnZwLv+lTF9ZKyr3mtB0+BezrT1GLUvR4VTn6LiDmHEHlOUC7m0i+4ezTix81DUoiQ4CRtowQZTsunAxNHIyVRXLiEWKByPld1LyF5jFmjKHy7X7YgdR5hIO7n5DK8670DHgfr6jKAmY1OqUZo/QPwvlhVMipcpnZHcIV98LpfTjKqW2NfBqG+a+AZSMlwxF9HomMPxDgMh2uEQHMyb4ku4xGHdnCOLyX4xXNdgEzhKYFjzvxVGaslQuY18jqYHV7rZ/wp+m8mmxYref7yLc9Dfv02lH9ty4bVMHFET6hUSkTHxKFq80EYN7QbOrUUDIsO/SfxvmN9o+vHD+sOtXXGnnAwhZwzv3vRrGKuP0DYwfN4d/U+v6h7rJmMZwN+h+aNaTID/Y2F943AzfrLJRdF+mOqcOpbnK7gxxP8JdZkNL3fG+t8N6Vb/KH14/ZY4bsLXDo0pP5vWmGu+z5RLmh6Fym1HyNbYYpv7lVuAAAgAElEQVTz7nSoMEJFox89dkvSqEjHU3EKqWpK/Fw4DlTCcIyHVK5nOnoHpr1tjYke0u+3slfgp2fNMTadDYFrfnuMuFIfU8yTnOiALNI0F2pNKoH1FcXFIQj7quOLwdFejfPjr0nmqu58oQGuTb2BN8dMczULcwM6PWuLy612I+GVWPb8ZupaD9xougocHUWbN07YjD1rMxscnU1LNCpMEtlvjDTFjKzD36cjZYwxhabJIzgOVPPYamYPadnSF/ejH+yWtZakUdH43/XdBc8DdSTvpySmYY1PwOsC0ZwkKHCcDIn1LsPmRgk+1amoyRUIKnwVnoHSUdKpMntE5DoPt0gK2BI/X8s5IdL1LFziBRqV3KYeVHadILf5/6gGVdJ6sT5Bl/TKa0zglaH9d5KFEz/zFhwayfOR7WzV/6oJyHKOjH4HDzeXT5bZiynkfxVDzrjg3c1HePazvprTp7eQO78YgNUFNn8yC3lYcDPBdyxh4eZkC5mmOywTLOQ3K84jaNMtib0WhwonB2aphZw0aj0cV9KxesZbyHK1B7g6h8BdLC7JCc9oC5kAtvLwg0JdJWcsNP8yi2JWC7OEh3wvKZ0NYDZDnSnkbCaQTzGcpDeheDZ5FZLfCAXhFSqrT+pDdi7ji682N4BSIcOaXBsz3Ic8MrwFn4qR2kz3vaJUmhnlQ15U5wwCb30aH7LaSoGJ7rugyWQfslKlRJcnrZGcqv0kPmS7ArngQ/Qk4JP4kFN+MFYL+hQ+5JQGV3ls1I97Qhl33ezPMWN9yFyqkLlMpiwEK/dFkKtyfonGIta/GXJXC+kzjbmsP+XPj5K+/xRLa4Y/kynkDIc0ez2QjnejTv6J1/O38AUDVEWKwHF0F5NBZmSUtW+/Cig4oAT/fA8na76oAOW3zogoa6VSgZ+iWiA2wTRj1Cwv/wyPstYDtGfCvQyPsqZTeqKIRcUmY2XD05kaZd1wS3U42qkQFi2UnczoKGvnPrVMvq2MjrJO+3A7tRKqJW0yLMo6uvppONsTzSmRP3hxeLcItkGrP0mUNc3DUdMHypQrfBS2yn0OFOpmOZ4GVdB6gdiHrM/cpVfO5j7mDOh/ksgUcvbSTP/Ho3kTnmCYfdyFe4j8bS+0WqEkoFKhgO9mioYV2rttt/Bi1Rkh+TslDVEqUeKgMSBCOMqkakk6/6FKZVJy7+nEswjc/4T3H9JmQGVjhapnjRuASwPPIujEGyE5AKVFdLVF/YvGSMXDrc4h6EaQwc9kn8cObc41AiXOoMVyfslDiHwaZRivS0FndEpT8Wh1ueOIfBVryJXr86UbOh+sZrj+96KHkBCRqBsfUKiBD5r6Gcn9s/MfQGpsiqGcZNkuBdBoQWmeA0uJTya47kVKYgofMEQVib78riBqTRA2INQW+uznKVqkdGn3//Uv5VG2fx5jfy5/vhwjNfJRNlj1hQnNa23eXVSNkY/Gpmf0emlabWp7IcpHrTuq5eRo98RY6pKeeaTUViHymh7AydHgjpFjHn78KR6OPy0Essk4fj5VzhvLJb7bdwtvl58QDoJlHGRyJYrso4IPQovzO4GYw+f40wmam9zaCsV2/AyiIL2LT0Hqxp2IP3+RHzh9Owq1GrZzheBBvq37HZrH13keLglYZmcP2U/GcojqzZOgffPYGBnt7I6kPsZEDS77BkMTHmh4nMLNF8ruKxD1LgUpGi28T3WFJvaNMVrItSRCq/5muN7nSlNokiN05RYBjVt1RBRLU27xeTVodck76I8jybMFolx+NtzvGVQB2pQUg2wTnXryAVz65hleAdpU2ixykMvliLMbjTjbzsZvNbkMuBQqZSkIN9lmNuIgRJRT83G1gowCFHJwK6Ceb+pDNvcZf6KfnycNtQhUmYVsEWL6uEHGXbmB2NuPEH34DOyXCgvfyw5TxcEmHIe824SqJbfqk4/H9H20gJc+IuwwiQpjHqxCka0VT/Tni9Jfa71ddL+1gxqVTnREyNkoXOh72CyTFuCQ1w61jzbBrYWPcGMO+S1NK/Z4VvRCj+N1saPPNdzdJCh7Q+M4FG1fELUWlMa2Zpfx8mKweTeq66pFrfniJKIDYkRUnpY7ayJ/dWfML3QYidFmBRA4Dv2u14NbXjvM8tkHTapZINPfuZTHhAlRz7957xNHbf9NdRmiq8a00HOvOEiNA74LFu5fRVHT5nFEWhl6vxZoVNsLSeTS5oB2TwSa0uESm/hTkLSNdHPDu0JC/PNVV4nCjEi3V7kkKOUHTeZJFp8oun8Y3x/YZZIoXzOnUCC331hwwSGImThLJHuVjRrq2ZNh8/wWYlctNaPYcZC7uQPDpwDX98L6pFCTW9/4jUO+YkhqNxY2V5ZAde+oebEqpJRoANs6QyA7Pgjy4L9E1aziSn6DuPwd4HWtA5D4WlTN6V25uUiwqwSfF9XBJQtlLY3vB4KLUEGQAnAPLA+5iILGISS34Gd3Cy0FOWcuPA6hnkK/U3xJnsdv0mRApI0Qqe6SXBIK255QWhEdypSrbnqTZf+Ux3pulvCQXyYK33B2b0whZ3cJ/YfxJQUG4QUtdpToYskspBy5gTerDooWVdqw59n2E56OP4D4K8/Eb+Q4lDryPW532YjkQInIXA6oeHoALtTcKEljovO/r6/1wt7i26kSvOSMWjxoj7X5aUGWKu7AYXhkF0y23cbXqjVvZOsPetMKs933SQbiEP1keEgzLKAgMIniFTIrBb5/1RQz+GpK4ver3ZUYer8Jppv7q3UD8a3gjG6HamCBp/T9Zb8rghrjimKh9x5JGlKXq3V5hb8qDxWXMN8NCcUfqP2bQj5UcpPoftJpDe904jdLNztskdxsUXGI4NWnEblL4NiaNNrwHBiG8ImrkPTYaJ0aruGAXJsmIG70JGjj6GRC9AA4LJ4NbuIQcKnSUdvyqcugmtsdinTq3CSO8IPjuhaQiRSe8C754ANQbqsuJVpwnByhTY/C80JVnWVuOj4NZIiofh6eD8vy9bBFo1fYILjgeXi8li4ewSlzIdT7D7gHl5Y8bk5S1kSM6wJDVLUZuEi2HYE49IJzYkmhGpTncciV+f7DX332vtVXPTc9spNEBq+0f+vmJKn0TGlpMlVgYsZm1PpUKDOF/KmQzeLnJtx7jLCdfyD+ziNBIU8chTc/bkRqrPH42rioChby7caLDcepaYdP62TpI9/heo2lkjQlurbCqQE4W3GNNI2JA5re64cNlOdagklCx4MtHrbHuvyU1EGaRtXndTssJN6vZDUoDt++TV8h0/hHhjXHfI+96dKwhoS0SFchayHDT6HN01XICqUMo942x3yqVyyh0K3sVPj2WWMs9JJW2M55bdB6x1fY8hWdHog/HF/KoPW1N65NvSxZfrH0kDKIfxyCkJOUa9lcoQLVDjbHg5EHkPRaguLGcah6sS8et1sEjZlvnn+STiEHdp8KmUZc+pLExSvkQaMkYp4FcdkvmQVubH8+G5hoeHS0PXUZ1HOEZCNSjVfIa1rwLgKp9q73Hnj415auNCbTIrTxSXidqypdyUvGIaneVVinQ3PSyoGQQtfh+Vqa5kRfKyUDcQ9JR2HLrRHmfgXO8SWlN5sKZ0Ran4VLko4GZV0dKod+kFtJV5/K4mXlP7/eWz3bEGWdXglGfTWojOwPYgr5P8uOPeAjEeCiohD7+CXezDbSnByW6I6s20+VPJbMu1U4sr7dgHxuojNr3kIOXHYNQVvMsxeRCuVQ8dS3eDDhAkL+EDYAaRv5Sts9/QY7qvsj/GGkxKxkaPGgHbaU9UdSlNiKksvk6Pa8FVbk90dqslgpkJ+679PGmMMHd4kVuo2HGgPv1cdvXnslF22Pki7ofPxrzPDYJ7khKNkmD5ouK4/pPG9YjE29WeVRoVdezPc0i/jWKTTDkbWUQuY4fBcsZM5alUfMe6bZ9NH5kbcVErsDCPv2j9vz9x8qtVl0vqC3kKn/fBUqfWcKv76eMf32fpP54g2BTiHT6UrIun1ipaJTyNy6rYi9Iv426HX2i2cBi6ZDG/RSJHs5HbH/sgS2ywdCGyuxYZDLkDhsPRw3doYsRbyZ5BTWiOm2BV4HGoCDuDyk1toBYXX3wPvSVxCdONNo7AoBtbZCdrW4FIsKSc71hHrHr6Qt4DiXUYiz65qOQjYeWesTgZgDEGmrO7JOMhaasHJdAYXaNDjuI5eCbHebh82sLPEhhyaOynZYSA2IWcgWIaYPG2RqeCQCxs1BakQU79O09nSB9SRB4b7uPANayg2sW5hpQbb9zBceM3rw/XcaLNYFBeneyXFwrPo58k0U/Fp/1lzOP1O/sNP9Pp0qwneAUATgbMW1vKVkXPg5fH21F2ytlbBSybDedxO4tFauzjrWz3Btvh1m6RI59H7VzkBzWupDVmYaPDjBOta32fyxsqmPeYTOx0vXLBBlv5JhiC4V5eWFj3By0v002NA8ZRgTKvh4D4z6C7fWPDc8X5d+19C/p8N5BJyizD/C+3mclEoMedMUNlYKrPnqGELuRxuUGvUrrJUY9FLIv72r6hFEUUCaYfgcVNYKdHssKOy9X/gjKTIhzfs52NhbodlNof9YpR1IjUs2mb+dpy2qHRfGf7naaqRqyd7XCx9wKeONsquaISlFi0dNF0LLKzVjv0P5AvCdLNz/tusvfNEOw/jI+q1VDo79hKC82MEjdbRw/QQ42LduDlkdXdKLsd+C0jrrG202FFOMqRStZnczcYHTZiNpxAbD9Q5rWvKBaPpGe693vYzZ0jwO1DbZMPCpLJsKqTKpeZ7/yuTToSelNLqGhGQN4hNT4fWoXBpsOEp/jaBCwjE+RVzbRK8wDp6OymUcgnwFH7FrtJC3Ou0fFqdUIMxNuN8JHSGLv5sWOr4+c5SNkGrTVdsA2uRAXnYyhS+s3DZCrsxt8r6c8IOreqbOQhbWCH0ua/3f0qf6OTxhtEXAxxSyRYjp/QdJ+XJpcQkbNhkyooWMl94Zvhn4G6xc7eA+xRhlm/Yt99qvhENed+SZLSz25u12Cz/YFfDAZ/PFASgUzfto+jW4VfQyKVZA+adj4lIQsO45bi+6BfcaXoac1mmf/2D5I9xa8hhF2xVG5SnF+fmkbVfnPMGd1U9QsnchfDFCzN08NeEe7m95hSojiptEOOufcXLILTw+/AZfjS+GEl3E/ro9/a7j1blgNJxVTlTogp6xrcN5hNyJx+C70kXmtzU7i3fP4tHnrhBBq6c5Rb5L5jmu2+qdQkJwMnrcqi+J7Z66J5Ecl4r2F6Wff7TOIf6+eselg3/ONdgLhdoaVfZK919vtglWXk4oubIJVAo5nB1UCI0SaFDUHndZCXVuR+SZKVjeaZtCLkPwwFmQ5feB20hT+hx/3Y17iNuyDXaFCwF9BT6ySbt8AdzxXZB9Xh5oa4xANlxz7Q9YXfYHin6N5DqmVcjoGvub26C45w9N8WaILWs6Ptr02d5eDDw/jITCXRH3mWmEOt1POaxtQs4jpmB/pHq3gJOdiq/mpG+uISNgHf8XYjzI8jVGQOv73SIGwir5Pt56i9OK0jXuUT2A1NcIczduBNLO3zmpHRSySIRbHZOUvUtKG8iV8Ui1OwobtRDZn5Oakw1FtVOUPvEU9O3T/xydYGSSZGc8mULOztL5iLFpk5IQFJsmg9GNe4hasglaLU+mgUIuh9MKIxXl3b5LiNlwHAIJipKGyJFro7GCSaTfeURsu8T305+N3EqFQjuNnL6n084i9OADXYXTvwsC2FihwtGehpHfGHgW4edfCMkAZDKonaxR7bQQFUztardDiPorjPc909Jj5+WIqkeagRJnONoqsa6SP6KfRPM0GrrGqZATGh02KrKd1f5A7CtddKwM8C7njQa7jTQnOgZP5Bdc4fn56+RBjVVGmtPawgeQEi/0038l2hRE1QWUp1homz7fjZRkDT9+iqIt1acIyk0oAeLAUuKTBe47eJqTYPrI8NW4EiiapprV+vx7oOFS+YAhirKts6QqcqepZrWrCAW6aXXzF3J8p23HymwRKGpEQ+LkqH9biJjWt7MV10Oj894qIEf1a90NfbQxevbzSWhlZBVzUMhkqHDGWIaOaE7BK45BQ8/mKxrJUWi3kR6iOXoJYRsOgzZ5dL9coUTRHZP5TQXPBfc/gPgjJww0LYVSBZv5M4yD818N/HlJcAMQPkprYOwiQ7/DgdlIffanUJObFI+tI+L7LDf0Ox0ZCbwlGpRwmqyy80J4W70bBvA80wNctDHQTObkC22zLfymj8bofaUNNIlBvNXN/0U4lUJo6WX88+l13g/qQJsUo3NVckghGpSvkWbl87YaNCnveD84IZji2MqUBhVVAZpUosAJn0+yc0/EKIw0KOcEChRL5bGVyWVIUg9FHHob5ufyrzSo/79qUCYf9//hD0wh5yChx1/9E7F/3ULsmQtQzxP4lxF9fhT7jLVauK4WFs5X7adI+pTzbB3H90sWn5DLUXjPMD5y906XTSIEFQ5qVDzYDXFn3+LC4OOiTFqOue3w5YFWeLXiEe4vvGpKQwIHz/I+aLCnCY58ew4PthoTuutfVLhdPt6y/qPNWYReCzatuAMO5UaU5qtF7ap5BO+evxNRXRpvqAPP6s5YX3Q/kmKNtZT553McOl9sALu8dvAruAtaXZk//btJMfd8IXB/V+XdLUrdSUdvfV8LltmaPDv4hThto/5eL4X7zSth8a/XAm2eCJbf0VIbRUFypBzq3RYsy1MV10r6hGteEzZEV6oblZdh/FoZKp0XirU/bjZXFFgsk2lReO9Ivj+o20RROAFxiL03/MwfTb8bPEL8fjngsGAOEBIC2VJyk5jOX2ZrA+2oBTwNSuY/XURTkrt4IbbbfNjd3wvlVUrKYdo4j6KIbjQTzn9Nh9VLsjJN/BdIzVcPiuo/gzvTF4qYmyIKXWzhvojz6QXvZ62AWFMaFE0quuBigQb1ujq0WlOfNm1JQ/JSpHwBuIeVB8x81oRJuIdwhC3lM6b7o2zv8P3/SoOKJxpUeyjt6kOeQ/3JooWD/YIMj3S4Bgwci0Mg6UUAAidO5ZWEeu5v/KIY/tM8iUAcDq6rpiNmw2lE+5vWcOUnzXHIs20cAkasR/JjoRqKSaO6x/4jcL3JeqTEGI/7DIs+x6HyuW9wvPJWcIlmCo8u0gJ1b3XF0TKbJfMRU3/HgN5Y7LZRkuZEWqTDs7ZYZ+Zv1r+frPmeL9qmG7UtV8jQ7WlrrPCVLg6hslWg56PmWJePknCIp2/v64A25+tjpUS1Jbo6d93caLDmS6zNu0O82eGA8qNKCRsGs9KUPPQAGh5rwm8IDpfeJBGkxaHBbeGo+PQXayVEw6Hm1V7CZqnzRokgLeDLc9/wNKfoPX+K7ieFX9h/OGIWbEb8nxKVvsDBe/0kxP4yBQiVCNAjv/KiOcCM7yBLEn8bNEFu4u+wXdJFMqKfvr3477fAeT35pSWiqinzXI99cN9fW4L3S/pXi+SOZ2F98EvJqGryYYd/fR6etypKFzmQqRBU8gI8X5WRpElxcjeE5D4O97B0ikfISyDCdXO6UdXJtn0Qh6Hp0qAibf15ha8vPmHttR9yK9P64ha3MLEBvzcCTCG/N1RZcyFlRQoJj4S7qxOfVSu9lvT0GcI2bQP9n1fIY8YhYvYaIEJcHpGe4bpyGt70mQdNrLgEHi2aebaNxeOW8yTLH9LzC/uPxOXqpkEuhrFxQKVz/XC8DFl40jSmVo97YlfBddI0KC3Q6F43bMxHtZSl7+/wrB3W590peT8dEXZ70TpdhUrj7xHQNl2FTPPo97oV1uXfIVmxh/p7vGiDVblJYYs1tkotR/fHLbEmL/GKxRJTu6vR8c8mkgqZrnYs44Iqc6rifBOqBS2+P1+PElCrE/Bi7QPJzVLFXe3wYsoRxN0Xb6Zo+13p3Dd42mM5NJGmiTB0ezEU8R+O0N6/QEPBfxKbMW+/iYgdMjJdTrn9wjmQTeorrVDJ9fHzCtgs6iDNKzYoZKr5LJ48fQ3R3ffCc39tPhOZeEcBhDY7Ds+zVSV5xfTIpAaXYX1ZWuFqZRxCSl6D58sykjQlrVyO0NzX01XIgJqnOekVqvn4NHIbxKjT79cqcyPaaglc4mlDQi6iilA5j4TcWgiaTK+lajT/uD5kzerF3vqhCDCF/KGIZeL1py/+hZG/LEV8gpA96ucRPdG+mbiYORcdjdh7DxC6fJVhdOo5C/h/R/Q1+oP1nbQou62axv/4qgP5k82OVcEh79ZxiNl7A8GrTkityfh83wg8n3ACISefitdEToYvz/XFlbYHEfMkQtwvAzq/6Iut5TYgJVxsRZESq3uzC3YX34XkJDHNSWElQ9v7beFXQLqesdrBGu1vN03XQnbI74TWp+piZe7dkjSoAl/7ou6mL7E2/04pfYiSA4qiwpgSWEmJPCSuqL+6KvLU88aavKYR4zwQHNBTd2S98/Ot4sxX4NDmUQf+0iOlxe4A2kzU11vIFddIWuA1dEfWl6svl3g+UOms4Ed+2IyyJpmJh+Pwuf8IpBy5iDA/MS+aDtR8/CYBe/cj9phEYBPHwX7RXCj9lkDzjI6MzRpF1k74HXZ+I8FFSSQakXOIH7QFzpu7AClSiUZkiOq+B54HG4JPkG7eVFYIaXgQXueqS1YCk6ucoG10HLJzpSRpTlp1cYQUXgfv12X4HNPmLdmpASKdZsItlGhQ4tcn2IzmaVDOcVQP2bSfNhNROpqTc5yQCMSk/Z3VLdJOONJOe+Stcv0NCtsKkMm8xC+k7HuBIWjUZTSObpmNXN7uktewX1oGAkwhZ1M5JSQm4+tW32Nw71bo0rouTl24iSHjF+Lw5lnI7eNhMmpNXBxejRoLbbxg7SqdnaEcP4n/d+SA8dCmCLmVqZEytvbxhP0UIXjnTcfp0FCBYCNTBTal88N9nHAs+qTFPHA8gdNI5XGpXxoe3wmBVZerrzShQdHvPutfGR5dBV7l8bIbTJQeLUr1bnblg7ZsrRXY8pmQvF/faMFvcLubIQXlls+2mdGggI7PjEFhomNrjkP3AGNg1Lr8lPrQ2Mjg7h5gzP9MNZnTgiOHHH0ChcjyP2fcxZ1lRitUCOwy3n+833W8OPTCOHZdfus+L1vBVq3ErobH8PZGmJHmREFzChm6PxN8zIfqHUBcQKxJv8pKjuZ3hPGfq7MX8aFprVgOVjbWqHlZ119rE1Jjk9Lcz8Ha0Q5Vjgs+6Gs1VkHLV/XQyQ4cbAq6o/RaYf5PW8yHJq1swcG1Qj7kmdoBickahPacJHwbaeTjULkE7AYK+Md/NwLkYjd8W+DgULc20EKgcckm9QPHyUz68bPxVMV2cWczK5tD/HdbDHg6rm9hcmhNroiY7sQVF5qnPxWPN+VRhTQ3Ri97nTUG99H1FJzGNbvGR+3HJ2n4Y2uTKGYZEFTyKv9su8Q1sA8RNrV8o8A2uQxBeYRNhlPceCgT9phQyCiLXKibwCumwhHyxCtpPi26X45IG8HHTNHUXMojExqUQqFCuLVQYcotuS40qUHC8+WuUHv/AZnCzTge3b86DZyMW/eETTFTyCJ4LO4XTCFnU5GRdTzwx3m4cWQFrKxU/Cgbd/2BV85dWhvpMLTIxCWkInz8BMislVCNEFvEdG/0iJmAowOcfjYWikg79eABC2Hl5QyXSdIZkwK6LoeygCt8JxuVof5+iuZ9teASbMv74vOJYgv+7e4APJl3Ga6186PERGOEs5ujFWITNHi06j6er/gLBToUQbGR5UU0p9tznuDZlvv4rGMxlJKgOd2ccA+P9j5H6YGFTSKc9eO7MOYWXh4KwJejSuMzCZoTKdaQS8GoPr2cSQS0/v4TXc8j9F4UOlxvIvm1UHBZzLN4dLzZ0MCr9HS2RnhMMlI1Whxsegrx4XFoe1H6/qONDoFL1qL+8caSz7/Y4hA0iUmodljgA5u3q20EJfXFTjrmFbfbrTZB5qZGyZViGhBd/aLHMii8XXmaE20Y3B2sERwlnMooFXKEDZoNztsZLj8ZI4TTviXpx/GQ580F1bffil/+8BpwYBNkuQuBay/+9uxuHwN3ZTu4PGWQUF+i/94uKG/tBFeoHmIqGqP39S+yv78e6gB/JOZrhthixghz6ifut93jeVC9OYSYfF2hzd8NdjZKQ6Urusb57STYRJ9BjMdQxLkZi5zon0/8Yut31xCT+yeTQhD6frfYgZBp7yPMUZoG5cT1gErzHGHKM5KyeS8alDYeKS5H+U2eOQ0qJCwKQSHhIMXMFLIkxBb1S6aQs6m4tvmfwtqtf+DgBmM1mu/GLkD+PD4YMcDIv9S8i0HwqvVAN4FPrJm3GPGPX+h8rwLXz7F2VaCNkDwjYtxSpL5+w1sWApUGcOvSBIp6lfn+N98thiZUX01J8N/mGd0JKF+Q//eLTouQEk+LNVVrEsCj42t9+6vhKqTqA7loc8/JUO5Uf0P/+Wp+0KQIFXHItiEq0FeXewmLp40Su4uuEKwyneVF6RK/vmJcaCkQjCx2gcdIiTsUqPeXkQp0oNg2kB9Q1wuFTI7Gd42biF2Ft4OTka0l1HhTKlVoed+oqA6W2EYoCnOTyXgaV/1rxv7DpShf9N9+ULL8wMHK2Qo1z7bjE59EnH+Li/0OGfze9AhbbwdUPihgH7ztEZ7NP8Pfq2NZwe4zZ5TeIFisLxZfRPDmWzqrS6AKuZTMjcJLBUUeNP0QIk4/MDnNcK1RFN4/CnzjoNG7EH0vgJ8XUcSoebcqC31JxIgfliLpTbBOdgLAXu0aQNFckH3U6NnQRoTz/9YzQ/OM7AGubEl+k5Ty8098vmrCnQLraP7WcxYaZKucNRBEu+MpauT/lAGp44xuFLtlXYSkNPqKuDLheFrf3Hc04+l5NHIefrkM4W0PGPq9/qgHLS97oV8uUyC4ESXjEBpl49If9vAjkCuQWv8yv+lLTtXA6zeUndkAABL6SURBVEFFvkqZcD8HhdIabwpfNNzv+bassKHSIcDJbBDqc9nQ7xJVylgggkSocESYvRAUaYfTUMUPNsyd3qBQ5UK41RFd/2ooU+cJRbh0H7dCWQAREDZTjpgJWfIGw9z4f6jKIBpCYhR7jICDy3TI5E6G8ej/ERwaidrthjGFLELG8n7BFHI2ldnKTQdw6OQV7FghHD1TI3+yva0NJo4ULIWE27cQd+Es4q5fg2raEv537waOFPkVaZFx1KXODO41XsLvycFzzWT+ft6nbObbotAZ3y1Cpq/HLeaIECOGbaG9wxC87RYCl14Q9VMyiTLH+4M4y4F705TX010pU8nR5m5/HGuzFxF/vRXdb+VohSrHO4MSXiQEif2K9nkcUWV/MxyuuFe3WTB9hFsFb1Txq4G9pXYjNUkc9V2gdX6eRnWoxDZoJTIyl5tSBT6t8uFo6c3CYpqmkeJuoKMhnazgJxG5y6HWdWFDcbGqUTkZHsFxqHKRAqDo+P93cVQ1jD7fuw0ptaU5jYpDiUOC++FB07mSPumi+wVu7FuiMZm/gTKtbRBKDEb0laDIcUIAILWkkVSK0fT9cjmg+vU34NofkP9BqT3Nxvd3yUrt2JWwvrAO8ut/iAubqFSI678e7mdGgwsS/Kdpm1xhh9A22+F1uge4uNfifhtPBNXaDJ+rDaFNkSh84lAKqL4GyisVoNWaxytwSHRqgSifn+EVVA4cfzxvKtw4lzF8CUW32NK0TzPtBocIJ+GI2jWxtLEkqe4q+jai7IR+p9RS4mgD6lfp+lPEua7pW4tWCZg4JpeEyqYZVOqmkKvrmoyDKWTRZ2Gxv2AKOZuK7n0s5NSIcMSdP8vPQF63KRAXh+ST5yQCfThYNW2ApHtPkPpInE+YdIxdq9qIPXod2ijpij2O7Wogcvc1cIlSFXsA185VEbT+unhRo8FRIFCvL/By5U3JQBtauIoOroj7C69JF2jnOOTrXw7PltySjDqmV3z2bWk8XnRXpDCpj2y5woNL4N5vxtSF5mIv9n0JPF50R/L5cis5CvYrjqeL6f0SUdVOSuTtUgwvlglpEM2bfV4HuDf+DK9XCv5B8+ZQ2hNOX+ZG4GoxDYlk49m2FFROVgj1uyQeHwd4dKuMlOhkRO//U1Khu3epgqRbT5DyRKzQ6Pn2bWoi+cxVaKLEEfk0W3XzutCePA4uWYLGRH7xBo2Bs/6Qmys0nSXK1WgJ1dWdkrKnHUzKl+1ge9eYIjMtPrT9SSjRFbaP10oXZ6A6zYV7wO7VSsl+Umqagv2hClwsXRhFxiHWbQDs3hlTeJq8X6ZAvP03sElcIhngx8nVSLDqDXXKYsn3axSeSJa3hVq7WDJqXCsvj2RUhloj1c8hUUGccxdDv8KmHeQKH6aQs+m6/V+HxRTyf0XwE92v9yHfPLoSKpWSf0uDTqPQvV19Ex/yJ3o9eyxDgCFgIQgwC9lCBPUew2QK+T1AyopL4hOS8EWj/vhhUCd0/pco66wYH3snQ4AhkPUIpKRq+KCuhp1H8/EmRHtSKdPPV5D1I2Yj+CcEmELOxt/HifM3QIFc+jZuaDd0aklUD9YYAgwBhgDwRaMBhjwFhIerswPO7jEG2jGMLAsBppCzubw0Gi2CQiPg6eZsOLrO5kNmw2MIMAQYAgyBj0CAKeSPAI3dwhBgCDAEGAIMgYxGgCnkjEY0k55HVCbi6/5TfutMGsp7vyYsIhp2tjawUVu91z10vb2dDdTW73f9ez00Ay563/zi5q8Kj4zhf+Xm4pgBo8iYRyQnpyAyOhae7s7SEe4Z85pMecqHzIVkGBEVw586OTnYZcr4MuIl72LjQXmrXZwcMuJx7BnZDAGmkLOZQN53OP5HLmDeiu04sd1Yv/V9783s614GBmPAD3MR8JqSUgCtG3+NCcN7pBt8cv7qHSxasxuv34QgMSkFVSoWx7Qx/XjlnNXtffOL68dJC/+qzQewfvthRES9g62NGlf/EGryZmWjDd3S9fuweM1ufhjke1w0bSjKFBcSwPxT0zMAlkwfhhpVyvzb5Z+8/0PncvHaXXw/fqHB9/pF2aIY+W0HlCxS4JOP9WNfQPnsf5iyHBRXQq108YJYOOV7vugMazkHAaaQLUyWpNz6jZyN129D4eXhYhEK+ZtRs3llOnVMPz4itH3/SZgwrDua1a8qQp92/2Xq9OFzeA/o1hwJiUlo2+9ntG1aA707SqeWzCwRfkh+cf2Y5izbhj2HzmJA9xZoVLsSklNS4O3hmllDTvc9N+48RtfBU+G38CeUKvoZflu1CweOX8SxrXMhN6vhnPYhD5++4u8jBZFdFPKHzuXS9XsIDYvC11XKIDExGb/MWwfaOC2dMSzL5ZLeAChR0Hb/U/BbOJY/Yfp2zDwUyOuDyaOl05lm24mwgf0jAkwhW9gHQgqLjnJPnLuBlZv2Z3uFHP0uDlWbDcKGRWNRrmRhHu2pC/wQFBKBhVOHiNDX072m/NAHrRpV5/t/mr4CCoUiyxef980vrp9UaHgUarYZirRzyS6fG20U7j8JwMrZo/ghUU7kWm2H8pnhihXOJzlMmk+HAZMw/Jv2mDR3HWZP+DZbWMgfM5e0E6TTpjHTfsdfx1dlWxcQbUob1PwC/boIhTsOn7qC4ROX4M7JNRbvasgufxPZYRxMIWcHKXzEGP44cRmzlm7J9gr56YtANO85Fqd2zoeHmzM/U78dR7D38HmTtKBpIZi7fBtWbT6IXh0b8cphxsKN+H3WyHQVxUfA91G3vE/2tLQPPn72Or4f/xs6tqiNR89ew9paheb1q6J5/a8+6v0ZeROlYXVxssfYIcZiIiVq9kzX6qXTgZ5DpqN6pdL86QXRbbKLQv7QuZjjSMr4yfPAdL/HjMT9Y59FeNPGjpQytXuPXqDdNxNxwX+xRfnAP3b+/y/3MYVsoZK2FIWsP05Mu3CQYlu2fm+6mwk6Uhz1y1LeT0ZlJ7/6oiRmTfg2yxee98kvnvZz2rjrGKb9toFXYEU+y4OHz15h0erd+HX8ADSpIxR0yKpGboQiBfOaFCqhRZ/ypJuPjY5zSelRIyVMR9rZSSF/yFzM8dZbx3RSUKViiawSxz++l3zkJWv1Mtks6Te6x7bOgY+XuCxjtpwIG9S/IsAU8r9ClD0vsBSFrF84Tu9aYAhA+ScLWX/EvXreD6hUrhjIZz74pwUoVCA35k6ULh2ZWRL6UAuZFPLWvSewb51QnIEaWWPkt5z/y+DMGrbke0jBUiDXT993NfSnZyHrj7PJj29no+avX7f9MGpWLctb+3qrLasm9CFzSTtGCh4kZf7z8B5o37xWVg3/vd5LG6CpY/qifo2K/PXMQn4v2CzuIqaQLU5kwoAtRSFL+ZAnz1uPkLBISR/y2cu3MeCHOTi/dxGcnez5uVKE8sLVu7M8OvlD84sbrj+2yhBRTsqDAtUWTxMqNGVVI7/rw6cveVcAtX/yIVMA14adxjKHdP2ClTvRtF4VNK1bhT/Gzsr2IXPRj1Pvg82O/n0pLMmH3LDWl+jbWSjFyXzIWfnFfbp3M4X86bD9JE+m46vUVA1fmpFoT4c3zeLrxmZnPnLfkbPgaG/H7/DNo6yJV9lr2Ez06dSYj0IODApD/Y4jMbBHC3zTtRkSkpIxcMw8ONjbZnkU7L/lFzefS0xsPOq0G44e7Rrg2x4tcOfhc3QeOJn323ZulbUpUI2RyWNRqthnWLByBw4ev2SIsr568wFmLt6MOT8PRL7cXqJvOTsdWf/bXNZuOwTy51NEOTWKX6BAwTGDO6N2tfKGuZFPnWhp2bGt2LgfO/af5qOsbW2seRohi7LOjpL6b2NiCvm/4Zfpd1PwSYteQm1ifSP60Iyfvsn0sbzvC5+/fMsvIETVotayYTVMHNGTT8oQHROHqs0HIW2ebtr9++04CqLYUKNjuu/6tM4WdKF/yi8uNRdzzisp4h8Gd87yDRRt7IjrvWz9Ph5jUkS/zxphiIQ/eeEG7yrYtWoyihTMk60V8r/NZdaSLSB3g57//cu89bwrwbxlZ2s5Lj6R9+OfuSSU+CTONLEUKKELazkHAaaQc44ss/1MqEwc8ZHtbN/PCiGajaODHaytVNlqbh+aX5yoajT37GiBJSYlIyIyBt6ebv/IP85WAkhnMDlpLunhTS6glJRUlhDEEj7IjxgjU8gfAdrH3EL+RI1Gg2Kf54ePpzExxMvAEJAFqc94RMdvr9+ESibN+Jj3snsYAgyB/08EaINy694zPHv5BknJKfD1ducDJcn9w1r2RIAp5EySC0WwUqteqRSWzRxheKueGnP31Fr+dxNnr8X2/aeg/zmThsdewxBgCOQgBK7ffowfp/1uyOhHVjWlbqWWnY/mc5AIPmoqTCF/FGwffhMp5IL5cuFpwBusW/AjKpYpwj/EXCFT4FBKamqWc24/fIbsDoYAQyA7IKCPmKekOrPGD+CDv6iR2+S3VTv5BD1D+7XNDkNlYzBDgCnkTPokSCH/MKgTdv9xFmq1NTYtHsenvDNXyGQdUyDQ3ImDMmlk7DUMAYZATkKA0ppu23cSf2yciby+4gh52vRTpPb4X1ejQF5vFC6QG5QgJSQ8Cgsmf8eMgSz8GJhCziTwSSFTEoa8vp58xDHxUCmxgrlCph3snkPnsn1KzEyCjb2GIcAQ+EAEmvf4Cb4+Hv9KEyRu8/3HAfzTaS1SKOSYPKoPnBwtpxzlB0KT7S9nCjmTRKRXyF1a1+Wr5cS8i8Pu1VOwZe8JPr2i3mfMFHImCYS9hiGQAxHQV0vr3q4BfyL3T40UMlEPyTigrG2sZT0CTCFnkgzSKmQKuOj23VTMHNsfRGNgCjmThMBewxD4P0CAkrZQVq9/K81ICpmSwlDqUNayBwJMIWeSHNIqZHrloJ/m88dFXdvUA6X+YxZyJgmCvYYhkMMRoBO4uPgE/gTu3yxkppCz18fAFHImycNcIVMWqtZ9xvNHRURHYAo5kwTBXsMQyOEIUJrN+St2YN6kwYZiFPopU3pXyntAldSYhZz9PgSmkDNJJuYKmV5LlX8oupEaU8iZJAj2GoZADkeAkoC06j0OAa+DMahnS3z1ZSk+KdH9xy/5sqdtmtTgaU9MIWe/D4Ep5EySCSlk86IClKWrUZfRJgp54epdPDXqxPZ5mTQy9hqGAEMgpyFAhU0WrtqJTbuPm0yt9lflMLBnSxBHuUP/SSheJD/zIWcj4TOFnI2EwYbCEGAIMAQyEgEqvBEaHo2k5GR4ubvAKpvlhc/IueaEZzGFnBOkyObAEGAIMAQYAhaPAFPIFi9CNgGGAEOAIcAQyAkIMIWcE6TI5sAQYAgwBBgCFo8AU8gWL0I2AYYAQ4AhwBDICQgwhZwTpMjmwBBgCDAEGAIWjwBTyBYvQjYBhgBDgCHAEMgJCDCFnBOkyObAEGAIMAQYAhaPAFPIFi9CNgGGAEOAIcAQyAkIMIWcE6TI5sAQYAgwBBgCFo8AU8gWL0I2AYYAQ4AhwBDICQgwhZwTpMjmwBBgCDAEGAIWjwBTyBYvQjYBhgBDgCHAEMgJCDCFnBOkyObAEGAIMAQYAhaPAFPIFi9CNgGGAEOAIcAQyAkIMIWcE6TI5sAQYAgwBBgCFo8AU8gWL0I2AYYAQ4AhwBDICQgwhZwTpMjmwBBgCDAEGAIWjwBTyBYvQjYBhgBDgCHAEMgJCDCFnBOkyObAEGAIMAQYAhaPAFPIFi9CNgGGAEOAIcAQyAkIMIWcE6TI5sAQYAgwBBgCFo8AU8gWL0I2AYYAQ4AhwBDICQgwhZwTpMjmwBBgCDAEGAIWjwBTyBYvQjYBhgBDgCHAEMgJCDCFnBOkyObAEGAIMAQYAhaPAFPIFi9CNgGGAEOAIcAQyAkIMIWcE6TI5sAQYAgwBBgCFo8AU8gWL0I2AYYAQ4AhwBDICQgwhZwTpMjmwBBgCDAEGAIWjwBTyBYvQjYBhgBDgCHAEMgJCDCFnBOkyObAEGAIMAQYAhaPAFPIFi9CNgGGAEOAIcAQyAkIMIWcE6TI5sAQYAgwBBgCFo8AU8gWL0I2AYYAQ4AhwBDICQgwhZwTpMjmwBBgCDAEGAIWjwBTyBYvQjYBhgBDgCHAEMgJCDCFnBOkyObAEGAIMAQYAhaPAFPIFi9CNgGGAEOAIcAQyAkIMIWcE6TI5sAQYAgwBBgCFo8AU8gWL0I2AYYAQ4AhwBDICQgwhZwTpMjmwBBgCDAEGAIWjwBTyBYvQjYBhgBDgCHAEMgJCDCFnBOkyObAEGAIMAQYAhaPAFPIFi9CNgGGAEOAIcAQyAkIMIWcE6TI5sAQYAgwBBgCFo8AU8gWL0I2AYYAQ4AhwBDICQgwhZwTpMjmwBBgCDAEGAIWjwBTyBYvQjYBhgBDgCHAEMgJCDCFnBOkyObAEGAIMAQYAhaPAFPIFi9CNgGGAEOAIcAQyAkIMIWcE6TI5sAQYAgwBBgCFo8AU8gWL0I2AYYAQ4AhwBDICQj8Dw5mTC1qR3uKAAAAAElFTkSuQmCC",
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAeYAAAFoCAYAAACCKVhnAAAAAXNSR0IArs4c6QAAIABJREFUeF7sXQd4VEXXfrdlk5BKIIQiCohiRRQUUARBQRBpCgoiUgSR3hSkSVek9y5NpIpSBQVEUcSK+tlQUUFqIL2Xzf7/mc3d7GY32ZIt927OPA9P2Ew7887NvvfMnKIyGo1GcGEEGAFGgBFgBBgBWSCgYmKWxT6wEIwAI8AIMAKMgECAiZkfBEaAEWAEGAFGQEYIMDHLaDNYFEaAEWAEGAFGgImZnwFGgBFgBBgBRkBGCDAxy2gzWBRGgBFgBBgBRoCJmZ8BRoARYAQYAUZARggwMctoM1gURoARYAQYAUaAiZmfAUaAEWAEGAFGQEYIMDHLaDNYFEaAEWAEGAFGgImZnwFGgBFgBBgBRkBGCDAxy2gzWBRGgBFgBBgBRoCJmZ8BRoARYAQYAUZARggwMctoM1gURoARYAQYAUaAiZmfAUaAEWAEGAFGQEYIMDHLaDNYFEaAEWAEGAFGgInZg88ApbZOTE5DaloGoiLDEBkeBrVaZTVDvsGA7OxcBOm0CArSeXB2HooRYAQYAUYgEBDwCTGfPXcJHV4YL/AaO7g7enVtY4Pdq9NX4sDRU+L3Px19GxqN2iV8t+85hsvxiRjR/2mrfgZDARau2YXaN1ZF57bNXBrT2cZZ2bnYsONDvL31Q2RmZVt1e7DRnej6ZAs89nBD8ft9H53EuFmr0f+59jayOjsft2MEGAFGgBEIXAR8Qsx//XMRHftMEChWjArHsZ0LoNNpzaheunIdjz07xvz5x6ProNVoXEL9+aGz8P3//sAvxzdY9cvLy8c9j72IFk3vwbJZI1wa09nGC1bvxNp3D4i1tWjaALfUroHzF6/i9M9/4bc/z6HZA3dh5ezRYrhT3/+KjTsOo3Xzhl57UXBWbm7HCDACjAAjID8EfErMoSHBQqN8c/wAPNm6qRmN+at2YN3Wg34nZjqKVqmsj54dbZn00kFk/O7yyQgJDrLqcvj4N4KMXx/1gqOhuJ4RYAQYAUaAEYBPibnPs22x59DniI4Mx54NMwUJpmdk4ZGnR+Kxh+/DH39fEBqmpcY8Z/k2fPvjGVy6el3c39aoWhmPNW+IPs+0RUx0hNjCmYs244NDXwjSf7hxffO2jh7YFW8t24YvvvkZ9FLQsP6toi4kWI/5UwaJ/9Od7+adH+HQJ1/j5zP/iPEfbnw3hvV7CuFhoaLNL2f+xdL17+PZji1xQ/VY7P/4JP78+wKaNLwTkREVQMfwzz/dGuOG9HD4SNH6Fq/bjW4dWuCRpg1E+/c/PIGPPv0Wg/t0wvsHT+CTk6eRlp4ltOoJw3siNT0TS9btxmenfkR2Tp7QwCeO6CU0dC6MACPACDACgYWAT4n55V4dERwcBDr6XTl7FJo9cDfeff+oINZda6Zi0ltv2xBzs05DBRnVv6MOIsIq4Jcz/+DC5Wu489Za2LJ8ojjytryfvq3ujeYdmv5qX4x/Y40gfCpSXViFEGxYOA6kIQ96baEgPCLuVs3uxRdf/0+8AND4W1dMFsZbJ776CQPHzhf9iVilQlr/8H5P4dFnRgtCpzGrVokp9Qn58ttf8OKYORg/rCee6/KoaCsdhUsdaZ6r1xLNcvx9/rJ46aDfn7twVfz/ha5t8Org7oH1NPJqGAFGgBFgBHyrMRMxk2bZtMNgNLqnHtbOewWtnx2D2jdWw9q5r+Dp/q/bEDMRYd3aNcx3zgUFRgyftBjHvjiNPetn4uZa1cU2unPHfPj41xg1ZTme6dgS4wZ3F1bSpEFPnbcRuw9+hiUzh6Plgw3MxEzzvNjjCTz+yP2IrRSN3Lx8VI2taJ6b6tu0uB931auFejfXxJ31apm1bulZK42YyUCMNOTKMVGCfDv3nSReQmjM8cOeQ6WKkcjIzMaTL7yGYH0QDr4zmx9hRoARYAQYgQBDwOca85C+nTFr8RZs2f0x6Gh7/bYPsXrOGJD1sj1iJryJjP85fwkXLtNxdqo46j164nthzEVGXe4S88vjFght+fDWOagaW6Tp0p3wgFfmYnCfzhj0QkczMY8Z+IyQuXi5ei0JMxZuEi8LxQtZifd8qrX57rk0Yv5g/QzUrVXDPASdINALwv5Nb6BWzarm3094cy0+OPQ5jr+3UJA4F0aAEWAEGIHAQcAvxEwWy22fGytQrHNjNfN9sz1iPnLiO0ydt0Ec6xYvkkbrLjG36f6K0EhLKk+3b46pY/qYiXnyyF5Cuy6pJCSl4uff/8GZs+eFRTaRPpUnWjXGW5MGiv+7QsxvLduKjTsPY+/GWQInqUgvNqQx31ijSuA8jbwSRoARYAQYAd8fZZPGTGXE5KX4+LNvMeu1/ujY5kHxu+LELJEY3f9Sv7tvq43qcZVx9PPvMGPhZvNRs7vE3KjtQHFkTORrr9x0Q5wwGJPumB0Rc/Exzv57ER16m9zEvj20WmjNrhDz3JXbxYlCcWKevWwrNu08LI6ymZj5r5gRYAQYgcBCwC8aM0H45z8X8PFn3+HF7u3MEbCKE7NEQOQDTJbIUiEr5omz17lEzJa+xNI4PYfMxOmf/8Q3H65CaIi+xJ0tjZjpTro0n+veI97ENz/8jr0bZqLOTdWZmAPr74dXwwgwAoyAxxHwGzHbW0lxYh4zbQU+PPYV1s1/FY3vvV10ISKkI94tu49YEfOwSYvFvfMnuxYitpL1vesdLXoLzbK4sdTide9h1eZ96Ne9HUa91M1KJAp6kpKWISyhSyPmDdsPCTerUQO6olpcJasxLl9NEBozaeU/HFkHnVbDxOzxR5gHZAQYAUYgsBCQNTFTmM1pCzYJV6QnH2sKiv1Bd86S+5PlHfOaLftF6E2y9m7TohHIIIvug8lqmtyT6Ai5S7uHcfstN+LSlQSMHthN+FC3fe5VcX9N/s+PPNgAFCnsf7//LUJnSi5Njoh5zopt4qmgMe645SbhEvbdT3+Y75invdIXTz3xsGjDR9mB9QfEq2EEGAFGwNMI+ISYpbtWsnAmS+eSiqQxS7GySTt+fc56YYEslbtvryNckXbs/QRLZw03B+kgkl20dpc50Ai1l46Paf5lGz4AReGiQnfW33y4Uvw//noy5q3ajv0ff2kl1gMNbhOxrGk+ClBCVtoUvatbh0es2p2/GI9d+4+LICHFDdTIYIvGaPnQveY+p777Ff1Gv4UJw59Hj86txO/phYJeLCzdv+j3UkS04lbZklHYh1tmo2Z1Nv7y9B8Fj8cIMAKMgD8R8Akxl3WBRJ7XEpIRUzECcZUrljocabyX4xOEzy8RsGWhCFqU+alKpWirWN3Uhl4C6OiZXLOqVI4WfsKulrT0TFxLTEGBoUAca5d2b+3q2NyeEWAEGAFGoHwgoAhiLh9bwatkBBgBRoARYATgG3cpBpoRYAQYAUaAEWAEnEOANWbncOJWjAAjwAgwAoyATxBgYvYJzDwJI8AIMAKMACPgHAJMzM7hxK0YAb8g8NOvZ3E9McVmbo1Gg+ZNilKc+kU4npQRYAS8ggATs1dg5UEZAc8gIAXOKT6apcufZ2biURgBRkAuCDAxy2UnWA5GwA4CRMxX45PwzrKJVrUUbKe0ULAMJiPACCgXASZm5e4dS14OECBivpaQgq3LJ9ldLfnOU2jZo59/L6LdUWCcVwd3F0F4uDACjIAyEWBiVua+sdTlBAEi5r/+uYgZY1+0WnG9m2+APigIPQZNR3JqOnp0eRQVI8Pxznsf4+/zl3Fs53yEh4WWE5R4mYxAYCHAxBxY+8mrCTAESrpj3rFqCuITkjBk/CK8u3wS6t9eR6yc4sh37jsRi6YPxaPN7gswNHg5jED5QICJuXzsM69SoQgQMVPSlZWzR1mtICoyDGu3HMCSt3eLDGhSMRgMgpzHDu6OXl3bKHTVLDYjUL4RYGIu3/vPq5c5AqXdMUvJT4qTNi3pxhpxqFk9VuarY/EYAUbAHgJMzPxcMAIyRqA0Yt5z+AuMf2ONTVYyWo7RaISKTLe5MAKMgOIQYGJW3JaxwOUJgdKIOSMzG0++8JrIhDZ2cA/cdEMc/v3vCvYc/hxPtm5qTolanvDitTICgYAAE3Mg7CKvIWARcOQuRRbYMxZswlenfzNjQHfOM8e9iFvr3BCwuPDCGIFARoCJOZB3l9dWbhDIzskVoTujI8NRIdQ6D3m5AYEXyggECAJMzAGykbwMRoARYAQYgcBAgIk5MPaRV8EIMAKMACMQIAgwMQfIRvIyGAFGgBFgBAIDASbmwNhHXgUjwAgwAoxAgCDAxBwgG8nLYAQYAUaAEQgMBJiYA2MfeRWMACPACDACAYIAE3OAbCQvI3ARyDcYoFapoVaXHskrNzcPSSnpiK0UxVG/Avdx4JWVAwSYmMvBJvMSlYtAVnYunnlpCgb0fBLtH2tidyEUfnPFpr1Ytv59UV8xKhxLZ40wZ5xS7upZckagfCLAxFw+951XrQAE5q7cjvXbPhSSzp7wUonEfPrnP9FzyExsXjIed9WrjcXrduPA0S9xZPt8h1q2AmBgERmBcocAE3O523JesFIQSE5JR3ZuLnoMmo5RA7qVSMzzVu7Ab3+dw9q5r4ilxV9PxiNPj8CuNVOtUkIqZd0sJyNQ3hFgYi7vTwCvX/YItOn+Cob27VIiMY+ZtgLRkWGYMPx581ruaNEby98YieZN6st+fSwgI8AIWCPAxMxPBCMgcwQcEfOAV+bi1jo1MXpgN/NKGrUdiCljeuOJVo1lvjoWjxFgBIojwMTMzwQjIHMEHBEzacxk8DV+WE/WmGW+lyweI+AMAkzMzqDEbRgBPyJgj5gLCoyAClCrVKA75jNnz2P1nDFCSr5j9uNm8dSMgAcQYGL2AIg8BCPgDQTIf9lYYET7Xq9hYK8OaP9oE+h0WpB7VE5egZhSr1Pjh1/+KrTKnoC7bquNRWt34eDRU2yV7Y1N4TEZAR8gwMTsA5B5CkbAHQRGTVmOw8e/tuq6f9MbuOmGOFxJyha/j4s25V5euv59rNy0V/w/NCQYq+eMRoM767ozLfdhBBgBPyPAxOznDeDpGQFXECgwGpGelS/+UQkL0Yp/dKSdnZOLxKRUxMXGsP+yK6ByW0ZAZggwMctsQ1gcRqA0BOhuWdKWpXakNTsK18moMgKMgHIQYGJWzl6xpOUcASLl5IxcZOea7pelEhykRlSFICbncv588PIDBwEm5sDZS15JgCOQl1+Aayk5dldZOVIPnVYd4Ajw8hiB8oEAE3P52GdepcIRIG35emoO8g1GuyvRalSoFKFnrVnh+8ziMwKEABMzPweMgMwRIPeorFwDktPzSpU0KkyHkCANp3yU+X6yeIyAIwSYmB0hxPWMgJ8RIGK+nGhyj3JUqlYMZmJ2BBLXMwIyR4CJWeYbxOKVbwToCDstKw8Z2QangKgQrEF4iI6PtJ1CixsxAvJEgIlZnvvCUjECAgFDgRFXC4OJOAtJlehgaNQqZ5tzO0aAEZAZAkzMMtsQFocRkBAgbTkpPdccftNZZChMZ3QYu085ixe3YwTkhgATs9x2hOVhBAoRyMkzICE11y08YiKCoNdp3OrLnRgBXyCw9YOj+OaH3zF/yuBSp0tJy8Cp737B+YvxqF2zGpo0vF2EnQ3kwsQcyLvLa1MsAhR681pyjjjKdqfQUXblKL0I1cmFEZAjAvNX7cDh49/g8NY5JYp39VoS+ox8E+cuXMUttWvgj78v4La6N2Ld/FcRGV5BjsvyiExMzB6BkQdhBDyHAJFyVnY+UjJN8bDdLZGhWoQEm+Joc2EE5IaAM8Q8d+V2bN/zCT5YPwPV4yrhzNn/0KXfJIwe2A19n20ntyV5TB4mZo9ByQMxAp5BwBX3KEczsvuUI4S43l8IEDG//+EJ9OvxBHYf+Axnz11CywcbYMqYPoiJjhBiNes0FC2aNsD0V/uaxXx53AJcS0jGrjVT/SW61+dlYvY6xDwBI+A8AmTwlZqZh8wc59yjHI0cqtcgIpTdpxzhxPW+R4CIed3Wg7ixRhV0bttMEPO+j07iwUZ3YvWcMcjNzUOD1v0xdnB39OraxizgwjW7sGX3EXzz4UrfC+2jGZmYfQQ0T8MIOINAvqEA8cn242E7099em9goPbQajqPtLn7czzsISBrzsZ0LoNNpxSTL1r+P5Rv34MiO+dCo1Xjk6RGYOqYPnm7f3CzE6nf2YdHa9/DDkXXQaQPTwJGJ2TvPHI/KCLiMAGnLiWm5yM23zh7l8kDFOgRp1agYzu5TZcWR+3sWAXt3zCe++gkDx87H5iUTUKtmHB7qOBSvj+6Nbk+2ME9OxE0E/tPRt6EJ0BdOJmbPPms8GiPgFgJ0r5yTVyCI2RuFiJn8m1VsCOYNeHlMNxCwR8yffvkjBr22AO8snYC7b6uDu1v1xcgBXfFijyfMM8xethX7Pz6JEx8scWNWZXRhYlbGPrGUAY4AETNF+HLTO8ohOhQIjCKCMTE7hIob+AgBe8T81rKt2LjzsCDdilHh6Nx3IqrFVcKyWSPMUj3d/3VRR/fQgVqYmAN1Z3ldikGA3KMys/ORWkb3KEcLjgjVIpTdpxzBxPU+QoCIeesHx0SAkRpVK+Ho599j1eZ9eKTpPXhr0kAhxbvvH8XMRZsxeWQv3N/gNhw58R3I+Gv+lEFo0+J+H0nq+2mYmH2POc/ICFgh4En3KEfQsvuUI4S43lcILFi9E2vfPSCieGVmmbKnPdy4Pt4cPwCREabgIfkGA2YvfVcQtFQGvdARg3p3CujTHyZmXz2FPA8jYAcBMvhKycgT+ZZ9UShfc2QFdp/yBdY8h3MI5OXl49LVBIRVCDH7LxfvmZWdiyvxCSLISFCQzrmBFdyKiVnBm8eiKx+BvPwCXEvxrHuUI1QqR+qh07L7lCOcuJ4R8BcCTMz+Qp7nLfcI0N1yQkoO8gzuxcN2F0CdRoWYSI6j7S5+3I8R8DYCTMzeRpjHZwTsIED3ytl5BUjyknuUI9Cjw4MQzO5TjmDiekbALwgwMfsFdp60vCNAxHwlMRu+1ZWLUKe0FnEV2X2qvD+HvH55IsDELM99YakCGAE6wk7PzEd6dtmyR5UVorBgLcJCOftUWXHk/oyApxFgYvY0ojweI+AAAbLEvpJkcg/xd4mLDoaaoo9wYQQYAdkgwMQsm61gQcoDAkTKyem54n5ZDoXumaPCOI62HPaCZWAEJASYmPlZYAR8iEBuvgHXU7wTD9vdZVSKDEJQgGbpcRcT7uddBFJyakMFFcjWgsK3k62FLz5H6M96d2EeGp2J2UNA8jCMgCME6G75enIO8r0VENuRACXUa9UqVIpi9yk34eNubiCQmlUHMFqYPgp29v7niNC/3ZDW912YmH2POc9YDhEgzSArx4DkjDxZrj6qgg4hek1AhzmUJfDlVKjU9EKNGUaTpuyjn+FhrDGX00eOl80I2EfgUkKWrKHhONqy3p6AEi4ttY5f1hMewcTsF+B5UkZAbgiQwVdqVh4ys30TD9vd9VcI1iA8hONou4sf93MegfSkOqDLZemO2XTJXHicTU4ChZ89XR8W9ZddIelvlObSaOQRqpaPsp1/lrhlgCOQlp4pstlER4Y7tVLKiJOXZzBnwimpk6HAlGtZCYVyNmscuE9dT0xBhdAQhAQHOVwSJR9ISk5FXGwMu2U5RKv8NMhIuLmQfYsMv0z/8+7nCjG2GjMR8pR5G8TcU8f0KXUTZi/bik07D1u1aXBnXbyzdIJHN4+J2aNw8mBKRIAIduyMVTj2xWkh/t2318GSGcNQqWKk3eVcvZaEGQs34dT3v4n6ejfXxPhhz+G2ujfatKc38aT0XOTIxD3K0f7odWpEl+A+df7iVQwcOx/nLlwVw3Rp9zAmj3oBuhIsuodOWGTGlBLbd3q8GUYP7OZIBK4vBwhkxNcR9gwmjdh3P0MrW2vMh49/jRkLNyMxOQ1Pt2/ukJjfXPou/rsUj1cHdTfvkl6vQ1zlih7dNSZmj8LJgykRAcoJu3PfcWxeMkFogS+PW4BaNati+qt97S7n1ekrkZyajmWzRkClVmHqvI24lpCElbNHW7WnL53c/AIkpMrLPcrRHsVEBEGv09g0G/DKXJGab+a4/iIFX7eXpooE9k+2bmp3yKVvv4/WLRqhZvVYnPruVwwevxDbVkzGXbfVdiQC1wc4AplXSGP2fQmNsybmzKwcpKZngHJDB+uDnCJm+tunnNHeLEzM3kSXx1YEAk/3fx1tWjRC/+faC3npLXrUlOX4+ZP1dq2Uew6ZiRtrVMHMcS+K9u9/eAJL3t6NYzsXWK2X3KPik3NAWrOSCh1lVy7mPpWSloGmTw4WR3Z0dEdl5qLNuBKfiCUzhzu1vJZdR+LZji0xoOeTTrXnRoGLQOalm33it1zcTzqk2p92QZ22YBMMBoNTxPzRp9+g8b23iyuvlg/di/vuvsXjG8XE7HFIeUClIdCo7UDMGNtPkDOVX//4F10HTMHJfcsQGV7BZjnHPv8eQycuRqtm96Jz22aYs3wb+j7bThyFSYVIOTM7H6mZ/o2H7e5eRIZqERJcFEf77L8X0aH3BBx/byEqx0SJYTfv+gh7Dn+BXWumOpyGjr/b9RyL5W+MRPMm9R225waBjUDWhbo+8Vs2o1hoWBZyg33jL2eJed9HJ/HvhSvQB+nw85l/cPTE95g/ZRDatLjfoxvGxOxROHkwpSFAb9R3PtLHijAkEjqyfR6qVomxWdLFK9fRf8wc3FL7Bnzxzc8I1uuwfsE43FyrurktjXs5URkGXyXtmaX71Omf/wSdFFi+rOzYdxwrN+2xOSkoPl5GZjZ6DpmBsAqh2LBwnGwsX5X2rAaSvFnnTBqzZOpV/Kfk1+zp+uAby6YxF9+DcbNWIzklzeYaq6x7xcRcVgS5v+IRII2ZjqVbN2/olMb8zEtT0bzpPRj0QkeQJffrczfgxFc/4cv9y6DVaMTRdU5+AfLz5REP290N0mrV0GvVwppaeln5dPcis1GcMxozWWUPn7RYHHlvWjweUZFh7orD/QIIgZx/6haG4SyJer3z++BaniXmhWt24buf/sDmJeM9ujtMzB6FkwdTIgJ0x/z4I/fjxR5PCPFLu2Mm7e/+dgOF1TbdL1H55cy/6PbSFHywfgbq1qohiJmMvvIUTsxBOg10WhXUKhXs3TFPX7AJ8deTSrxjTk3PxLCJi5GVlYNVb41mUlbiH4eXZM75q67ZGltSnIuss02n3NJnT9YH1fnDakUGQwEKCgowY9Fm5OcbMGV0b2g0GvEySi/dfUbORr/u7dC25QOiHxmJdWjdFDVrxOHM2fPoM2K2+N546XnP2k0wMXvpweNhlYPAmi37sWv/p8IqOzREL1yCLK2yN+w4JO6SpLfiNt1fQa2acZg9cSBCg/Wgt+ZPTp7G3o2zoFapkZiei5hwPS4nyjvSV2k7RF9MsZF6kQlLyj714pg5iAirIE4X7Fllk8FctbgYjBn4DMja9dmBU4Vf+IKpQ4Q1NxW1Wo2qsZ51LVHOk8aSSgjk/kEGU5ZGkdKxttTCO5+DbrHWmHfs/QRT52+02hjyxiBXwJTUDDTtMBgTRzyP7p1aiTZ0WkZ3y1Lp9PhDmDSyl7Do9mRhYvYkmjyWIhEgLXjMtBX47NSPQv47b60ltMDYSiYjJzLuovvUbz5cKT7/9uc5rNi0R5B1aEgwGta/VRxr31mvlvBXTkzLFRG0oDIiTaHGX0TGImJZZh4qhpP7lBr//ndFvLRcuHxN4EBfSqRh6HRa8blz34nihWb+lMEgX2+ywi5eyJ/5xAdLFPmcsNCeQyD3tyKN2Zd+zLp61hqzOysiTTopJQ2VY6KdCrLjzhxMzO6gxn0CEgE6rs3Lyy8xsEjxRROh0/FXZITJcpsMvq4kZZuT5FAUrespOaDIX0oqQVq1IOOryTnmkIlx0cFm1zEiXdKAK4QGK2lZLKuMEMj7xfMuRs4sT3dH2YnZmXnK2oaJuawIcn9GAAC5R2Vk5SMtq8g9KjhIg1C9RmjQSiqVIvXC1Sszpyi2d3iIFhVCitynlLQellV+COT/j4jZIla2ENH7n7V3nZEfGHYkYmJWxDaxkHJHgIj5ih33KIqilZ5lQE6evBNYSPiGBGkEAZOmX7zEVQwWhmBcGIGyImD48VbTaYw5X0UhKXv5s/Ye1pjLunfcnxFQBAJ0F0t5lrNzbclXq6FjYZ2IACb3QpwbGxWMpLRcYVVevNAJAOVtJsMwLoxAWRAwfG/SmMkArKR8zN6oV9/LGnNZ9o37MgKKQYDcoq7Z0TClBURW0CHfYERGtryjgIWH6qDVqAQxl1QqR+qh08ojNZ5iHhAW1AaBgm9vLXJULgkfkfqxFPDcqFc3ZGLmx5ERCHgESFu+npojiLe0Ui0mBJcS5Os+RfGxY6P0QrMvzViNiLtShJ615oB/sr27wIKvbrXIxyzlYS7KMlXoyGzOPuWpz+r7f/fuwjw0Ot8xewhIHqb8IUB3ZHR8nZSe53DxFYK1QhtNyXDc1uFgXmgQHR4kXi7SMh3LFx2mAx1rk5sLF0bAHQQKvrzNopsbqq/VpM73VzcxpWqVe2FilvsOsXyyRcDVeNh0f5uY5li79vWCyT2KiDk+ucjVy5EMlnG0HbXlekagOAKGL24r8W7Z7NdMVtr27qCl/M1u1Ksf/FURm8HErIhtYiHlhgBZYadn5iPdhXtjynEcFqKRXX5mco8iV68sO8ZrJeEeFqxFWCi7T8ntuVSKPAUn7jDHypZkLq73euOzptkvioCIiVkR28RCyg0BulumYCKuFgrcQf7B9iy4XR3LE+1Dg7XC19qee5Sj8SnoCFtoO0KJ6+0hYDh+hzkWtklDFlfO5jtlUx7lojtnT9Wrm/+siA1hYlbENrGQckKASDkpPVeE33S1kJEVaahX3SB1V+dy1J6++MjgqyT3KEf9KUxndFgQk7MjoLjeBoGCT+4qipQtkbKF6ixI2gufNY/8TxG7wcSsiG2lrLfBAAAgAElEQVRiIeWEQG6eAddT3Y/mFRGqRYFRhfQsx4ZW3lx3RKgO5JJMPtjulkoRQaAsVFwYAVcQMBy5q5hGXFxD9s5ndaufXBHTb22ZmP0GPU+sRATobvmaA5ciZ9ZVtWIIriRlmeNqO9PHk20o8AmRanxKjkhW4W6hE4DKUXqOCOYugOW0n+Gj+n5Zuaa1KVGN3AsTs9x3iOWTDQJEytk5hjJpmNJi6G43SKPyyFjuABQToRdhQtMtYnu7Mw71oWhgIXp2n3IXv/LYz3DoHvOdshRFxBQHrMjkS7pz9mS9+vEfFAE3E7MitomFlAMCrrpHOZKZomjRMTJFDvNlobthikZG0crIqMYThd2nPIFi+RnDcLCBXxaraXfaL/O6OikTs6uIcftyiYCUm9gy41JZgQjSqRERohORw3xZyOCLAp24Y7xWkpxk2S3urDmOti+3UrFzGfY1kMywi0UAk0Jom820PVqvaf+9IjBjYlbENrGQ/kbAYDDiarLr7lGO5Car5uw8A7IsUiw66lOWeopARlG7ErzwMlAlKhgaDUcDK8v+lJe+hr33FcXKlk6vffBT0/E7RUDMxKyIbWIh/YkAxY5216XIkdyURjE2Wm83ZaSjvq7WkzZL2vL1lFzkGzx/fE4RxMhPm7VmV3em/LXPf79hMb/l4n7M3vms6fStIsBmYlbENrGQ/kKA7pVz8wqQUErGpbLKRlmd6LI3zQOGWKXJQkZapJR4M1435Z8mguY42mV9KgK7v+G9RuYFeiPCl+W5jeX4mqe+UQSwTMyK2CYW0l8IkCV2fFI2yuBR5JToFEWrrK5LpU1E7lFEmpQ9il42vFXUalNOZzoJ4MIIlIRA/s5GhbGypazLUl5m737WdP1aEZvCxKyIbWIh/YEAkXJmdj5SM72fR5ncjYJ1GhFRzBuFSDkrt0Csx9uFAqiQOxiTs7eRVu74+dsfEMKbXKRsf0or83S99pmvFAEaE7MitomF9AcCnnaPcrQGynOcmpUnjs49WUKCNAgP1Qpt2VeF3ad8hbQy58l/t7F/7pi7n1IEYEzMitgmFtLXCJB7FN3FupJxqawy6rRqEayD/Is9VehEmY6Wk92M7e2uHHQCEMnuU+7CF/D98rc0sQiG7TvVWdvzS0Vgy8SsiG1iIX2NAFkt+1LDlNZHxJxrMB2he6KEhWiFMVaiF43XSpKTLMDpbpsLI1AcgfxNTUs8xzZF/CqZrMtSr33+pCI2g4lZEdvEQvoSAdKWE9JykJfvPSOpktZDXzpx0SG4nJhV5iUL96hIvdDAyeXL10WnVSEmXM/uU74GXgHz5W940C9Sant/4Zd5XZ2UidlVxLh9QCNA98oUEcsfGqYEbFgIZX0yltnoLCosSCSoSM10P3tUWTeb/JopBCi7T5UVycDqn/f2Q0XZpSQDMCn/shc/a/ucUASQTMyK2CYW0lcIEDFfScr2WAxpd+WuEh2M62XQdKVgH1e97B7laH2mE4BgJmZHQJWz+ry1DxfSr2SVLdlfe/ez7kUm5nL2qPFylY4AuUdlZOV7PdCHMzhR2EyKP+2u5l4pUi/uqT0Z29sZue21CQ/RokIIu0+5i18g9std/bBVrGzxJixdHnvxZ1D/TxUBJ2vMitgmFtIXCNCxL2nLcinke5yeZRDpGV0p5B5FREgat1xKXEUOOiKXvZCDHDmrWvglVrZ+4HE5LN+hDEzMDiHiBuUBASLl5IxcZOd61oe4LNhpNSpUDNcj3oXkGXSXS9bQ3ort7e56goPIFYzjaLuLX6D1y17+SNEds3S37IOf+pePKQJKJmZFbBML6W0EKCeyJ/2HPSUvpVKkI/Z0J+NoU3vK8ETELLdC+afJV5sLI5C9rFURCNIxthQDzPy5sIkH64MHMzHz08cIKAIB0pYpJ3K+wfcuRc4AVC0mBJcSHLtPadQqEPn5yz3K0VroBICim3H2KUdIBX591uJWJd4p06mPiOdewl1zWepDhh5RBLisMStim1hIbyFAXwAU3Ss53X8uRY7WRjmUidQcZYWKDg8Caf7OateO5vVGfVSYDnQHzu5T3kBXOWNmLnpMCFtki22S3dufQ4d/rAiQmJgVsU0spLcQ8HU8bHfXQffGiWl5JeZR1us0INKj+2gvJo9yV3yrfhxH2yMwKnqQjAWtTdmljEbTXTMKs0t5+XPoyMOKwI2JWRHbxEJeT0xBhdAQhAQHOQ1GRmY2UtMzUaVStN3jUzrCTsvKQ0a2a1bPTgvgwYZEvGEhGiSk2r87JuJOy8z3aWxvd5cXFqxFWGjp7lO5uXlISklHbKUo1q7dBVrG/TLmPW4hne8yMlcYfUjGqBSJxsSsiG0qv0Kev3gVA8fOx7kLVwUIXdo9jMmjXoBOqykRlE+//BGzl71r7vP+2zNwS+0aNu0pTOVVGblHOdplOqrOzjHYkC+lWAwJUpdI2o7G9Uc9BVChO/HihTSoFZv2Ytn690VVxahwLJ01AvVvr+NQTNr3Qa8twPI3RqJ5k/oO23MD/yGQPufxwnzMUh5m3/ys8MqH/lu0CzMzMbsAFjf1PQIDXpmLsAohmDmuP67EJ6DbS1MxeWQvPNm6qV1hjp/8AYPHL0T/59qjY5sHER0ZDr0+yEbTJm2Zch9T+E2lFCIyChxi+TIhuUeRJk2JN5RSKExndJit+9Tpn/9EzyEzsXnJeNxVrzYWr9uNA0e/xJHt80s1Gjtz9j/RLzMrm4lZAQ9B+ux2dvMwl5R/2VO/Dx97UAHoAEzMitim8ilkSloGmj45GO8snYAGd9YVIMxctBlX4hOxZOZwu9pWl36TcOvNNfHm+AGlgkZBO0o6FpYz2pRXGUaVOIKnEllBJwxmkjPka7xWEp6VIoIQpLM++Zi3cgd+++sc1s59RXSLv56MR54egV1rpuK2ujfaHepaQjKeGTgVowZ0w9T5GzF38susMcv5IQaQ9sYTwupa3DFLJl8++Bw+br/MkTGJx8SsiG0qn0Ke/fciOvSegOPvLUTlmCgBwuZdH2HP4S/EF3XxkpichmadhqLlgw2Ql5+PjMwcNLnvdvTt3g7B+qK7afILvpbsn4xLntjJqhVDcCUpCxq1GpUig0R6SjoBUFoR7l1ReqjJLaawjJm2AtGRYZgw/Hnz7+5o0btELTgrOxe9h7+BZg/cjSF9O6NR24FMzAp4EFJnPmkhpe/umCMm7FMAOkzMitik8iqkdKx5ct8yRIZXEDDs2HccKzftwbGdC2xg+e3Pc3i6/+vo2r4Fmja6E6lpGZi9bCueaNUYU8b0Fu2JlLOy85GS6Zl8x/7Ym1C9FkE6lSDm7FwDMjyUu9kfayGNP0SvMZMzXV3cWqcmRg/sZhaHyJb2j/bRstDLCBE5FdKSyT+aidkfu+j6nCnTOxRzUy60zi7Rfdkz9RET97ourB96sMbsB9B5SucQkDTmT3cvQqWKkaJTaRqzRMwnPlgijIao7D74Gd5Y8i6+PrjCHAJQzn6+ziEDkG8znQBS0g2ll7AQrdnymoiW9m78sJ4ONWbpmPvp9s1RISRYtN+48zBaNL0HHVo/iDYtGikdmoCVP2VqR/trK+7IXLxVGesjX9+jCEyZmBWxTeVTSHt3zNMXbEL89SS7d8xS+63LJ+HuQiveHXs/EfeO/zu2XoBIxlLphfezSkaVklSQ9YyStWXCn46xyaqc7hpJ46U75jNnz2P1nDFie0q7YyZDr3fesw4YsWjte2j/WBO0f7SJON7mIk8Ekl/v5JNsUsWzVkVNMVn7y70wMct9h8q5fC+OmYOIsAqYOe5FG6vstPRM9Bk5G/26t0Pblg8IpAaOnSfuWxdOG4Lrial4ZdoKVK0SIz4bDAXIyCmAWmVEqoKPsknDJL9mQ0EBcvOMyMxRrtZMLmAGg1G4e2k0ahRZZU/AXbfVxqK1u3Dw6CmzVfaGHYdw9MT3wmrbXuGjbGV8YSRP7mIKhCM04KJjam9/jp6+WxEAMTErYpvKr5D/nL8s/JgvXL4mQOj0+EOYMro3dDotUlIz0LTDYEwc8Ty6dzIFxad2IyYvBR1rU3mgwW14a9JAVIyKELmNc/MLQD60lBKR/JiVVkirpGAi11NyxQtIbLQeVxLlk6rSFTyDtGoQMZPxmk5k0goS945L17+PlZtMd4GhIcFYPWe02Sp/zvJtws7gmw9XMjG7ArbM2iZOeMoX6Zdt0jxHz3hPZkjYF4eJWRHbxEJevZYk/JkrhJruEh0VOgLVajXivpKOSclfmYiZSnCQBqF6jfmzo7HkVB9VQSf8P6W42eGhOqFxpCnwrpl8sumOnGKVUyFiJv9mum7IzslFYlIq4mJjOOmFnB5AD8mSOP5pD43k2jAVZ+1yrYOfWjMx+wl4ntZ3CBAxU1AOSwU5JjwI6TkG5BSSgu+kcX8mrUaNmAiThimy7xSWuOhgxKcoy2WKXozobplOLqRCgcBio4Ot3KfcR4t7yhmBhHFdxTm2OUa2FCu78KcwoPBCfcybO+QMi1k2JmZFbBML6S4C5B6VmZ1vc6dM2ZoqhutF0gelFCLlrNwCsR7LQu5GwTqNiGSmhEIacRWRlMN0tWBZIkK1grAtfZuVsCaW0TUEro99xsS9dMVsJmNpDMlnyvP1ld7a7pqgfmrNxOwn4Hla3yBQWvaoiFCd8GtWgvsUpUqkqF+kLdsrMRF6YW2uhBCjhDvdlSeX8CLB2ad887fhz1nixzxjdl+U8iubfkrG2pJBWPGfZauvPGebP5ft9NxMzE5DxQ2VhgAZR9FdrHSHaU/+ajEhuJSQJeulkUFUbFSwyBlNoUTtFZ1WDbp/vmZxNCzHRYloX5H6Uo/e6SWEAo8QeXMJTATiR3cvMVa2hd5szs9sGSu7LPVV5m1VBKBMzIrYJhbSHQTy8gscEhUF6qBjbcmYyp15vN2H3KOIeJMKjddKmo/ILN9glLVvMxl40fG1o1MKCtWp06i9DS2P7ycErozsUaQxS3fJFrGyxfG2Fz5Xmb/FTyt2bVomZtfw4tYKQYCOqBNScpBncOwSRe5HiWl5sszOJGmYpAk7cu8irSIuJgSXZXoCQBbXUWFkvJZt8mEtpZD7FB3Ps9askD84F8W8MqJnoQGjSReWSLjQsdlrn6suYmJ2cau4OSPgGQToXjk7r8ChhinNRsE6wkI0ssw2JQXgSM10LnsUaddkOOVse88g7two9AJEgV0ovrczxdJ9ypn23EY5CFwa2rOEO+aS7pY98/uqizcrAiTWmBWxTSykKwgQMVPQDce6ctGoRIDZOYZS76NdkcETbS0DcFi6Rzkau0pUMK6n5YiIWnIpZGlNd8cJqfaN1+zJSXfr5ApG2hSXwELg4pBefllQ9aWb/DKvq5MyMbuKGLeXNQLCyjozH+kuZlyiI2MKeEH+znIpxQNwOCuX3AKoELGStkz5r/MN1u5RjtZEJwDSKYCjtlyvHAQuDHrB4g650GtZulNWiZg5XqmvvmyDIkBiYlbENrGQziJAlthX3CRXckeCUYU0GSS5sBeAw1kMqB0dA2fmGJw+NnZlbFfbklEa6bzJGc4dxxcfn7Rmvmt2FXV5t//v5T7iblkUEUtEYmOLz16ov2ElE7O8nwyWLuAQIFIm31i6X3a3kA/tlSTryFrujuVuv9ICcDg7plwCqEjRysh4jfbHnRIcRK5gQUzO7oAn0z7nB/SxuWOWyNnar9nCMKyQvMtSX3PV2zJFxFos1pgVsU0spDMI5OYbRHKHspRQvRZBOgp+4Z52V5a5pb6OAnA4O4ccAqiQZTX5kRePVubsGqR25PtMLmNcAgOBc/37+WUhN65Z55d5XZ2UidlVxLi9LBGgu2WKu0x+vGUtRAJ07Ep+0L4uzgTgcEUmfwZQIfcoOsYuKVqZK+ugE4BK7D7lCmSybvtvvxclzyif/rxp7VpZ4yIJx8SsiG1iIUtDgCyWs3IMbt9hFh+brKEjKuisEiz4agecDcDhrDz+CqBCp470gkOBWzwVJjQqTCcsu9lK29ndl2+7f/r2LxJOysss/caLn2u9vUa+oFhIxsSsiG1iIR0R82UP5yQmEsjNMyIzxzphhDd3wpUAHK7IQVG0ktPykOeiRbQrcxRvS5bU5B/uinuUM/NxHG1nUJJ/m7O9B0BlEXCzKJGFdVap4lmmyvq59oZV8geHbOGMrjhIKmJJLGR5QoAMisiKOiPbuaAVzmJDQTpio/XCH9pXxdUAHM7KRYQfFqLzOEmWND9ZUMdG6nHdDfcoR2uqEKxBeAjH0XaEk9zrz/YaKJlfW4jqRVVZzGJEnU1MzHJ/Nli+AECAwlR6y/c4PFQnHCrTsryvNbsTgMOV7fNlABVKpkEG2N6KPlYlOhh0F89FuQj81fNlcbdMeqG9rFJShgtP19+8eYUiQGONWRHbxELaQ4C0ZcpB7Kk7THtzkA9tfBlcfZzZOdIw6T7WnQAczoxPbXwVQMXdaGXOroPa0QlAdBi7T7mCmdza/vncILNIxfXk4rJ6sr7uluVyg8KuPEzMitgmFtIeApQCkcjMmyVEr0GwTiNeALxVSMOk4m4ADmflCg/RikAOaU7G3XZ2XMt2FK2MApuU1T3K0dwxEUHiDpuLMhE4032wuGOW7pZ99fOWrUsVARgTsyK2iYUsjgC5R11LdpxxyRPIkS9uepbnrIstZaIAHJUiKeOS+wE4XFljXMVgxCflgPDzdCGL6QohWp9Yswu3sii9SNjBRXkInHlmqImUvZDaURx/l5BKst52JmblPS0ssSIQIFLJys5HSqb3734JEMoLHBWuEy8Cni5E+pRtKcPF2N7uyuGtACrEj7FRwSKjF+Vb9kWJDNUiJNiUTYuLshD4reuwwkBfIii2MLY2Bfby7ud6OxbbBYquxWhujUxygLPGrKznmaUl20qjEZ52j3IELAXKoOAlniRQTwbgcCS/ZT0dN6dm5HmUQMlQjoKAEDH7srD7lC/R9txcvz493HODuTDS7bsW2bSm75Mp80wxtKeOoRje/i9MzP7fA5bABQTozZasfeke05eFdLK4mBBcTsjyyLTeCMDhrGCeDqBCx8rk6kXH8WQl78tCyT6kEKa+nJfnKhsCP3cZ4Zd8zHe8t8BK8MPHv8aMhZuRmJyGp9s3Z2Iu27Zy7/KKAKUN9ESIR3fwk9IPesINyFsBOJxdlycDqJArFp0meNOorLR10UsB3dVzUQ4C/+s8yi/C3vX+fKt5M7NykJqegQWrdyJYH8TE7Jdd4UkVjQBpywlpuX6JYS0BR/eoiWlli8ntzQAczm4wGd3ERevLfCVQ5B6VLXLo+qOQDBTKlFND+gN99+b8seNov2jMd38w167A0xZsgsFgYGJ2bzu5V3lFgO6ByF850cd3mMXxDg7SgI5PyyKHtwNwOPuMUAQtqEjTdd+Iju6rM7LyRQYpfxYiZrqz5zja/twF5+f+scMrUrZlEcRHWH5Jpcj+y/QbD9bfs3cOE7Pz28QtGYHSESBivpLkP63MUjoiAbrjJmtqV4svAnC4IhNF0aKsXO7cDVO0MnpJof7+LhQIjNbCxOzvnXBu/tPtX7VqaB1EhEja+vjFU/UN9r/FxOzcFnErRqB0BMg9irQyX4TGdGYvyPq4YjgZO7keR9tXATicWQe1Id/jYL3GZWtqIkC62/Wle5SjNUWEakEvC+w+5Qgp/9d//8Q4i3CcdvyOJf/mkvyc3ay/7+BsJmb/bz9LEAgIEDH7MpmEM5iRJTDJle5CHG1fBuBwZg1SG4qilZ5lAEVSc7YIS2iV96OVOSuP1I7dp1xFzD/tv237mim3lMUxtS8+3/fhG1YLNhgKUFBQgBmLNiM/34Apo3tDo9H43V6B3aX881zyrE4iQAZfFKrSnWNjJ6dwu1nVmBBcScgqduhmfzh/BOBwdmFk0VwxXOe0tbuIuhWp93oMcWflt2xHLz/kc86GYO6g57s+37QZ77vJLGZqdHiW1bw79n6CqfM3Wv1u+qt90aXdw36RT5qUidmv8PPkjhDIyy/ANRncYdqTs0KwVgTVSMnIc7QM+CsAh0PBChu4EkCF7tgpupcrpwXOyuGJdvTSoNOy+5QnsPTWGF89NsEvVtn3fzTDW0vy6LhMzB6FkwfzJAKkLV9PLZtrkiflsTcWxWtOTstDnqHkMJT+DMDhyvqrxYTgkoMAKlK0MnpZ8pd7lKM16TQqxERyHG1HOPmz/tSjk/wyfeMj0/0yr6uTMjG7ihi39wkCdPdEx9dJ6Y61UZ8IVMIkRFRhITokpJZsmezvABzO4uNMABUy+KITAm+m2nRW3tLaEebB7D7lCSi9MsbJlpP9ojE3OTrVK+vx9KBMzJ5GtJyNl5ubh6SUdMRWivKoq4o/4mG7u3WUGzg7z4AsO2FC5RCAw5V1lRZAhSye6Q63tJcQV+byZlsRQrVi6e5TdCITn5CEShUjodU4n0IyISlViB4THeHNJQT02CcfmWLlxyyyTJlXLFlpF/6CskV5qP7BT6YoAlcmZkVsk/yEJOJcsWkvlq1/XwhXMSocS2eNQP3b6zgU9tMvf8Sg1xZg+Rsj0bxJfZv2wto5Mx/pPsq45FBgBw2kSF7kZ128yCUAh7Nr1AdpEKbXiAhrloXWSHe3lP+awqIqoYQFaxEWat99ip7BMdNWIDPLtGevj+6Nbk+2KHFZROLrth7App2HRVzl0JBgfPPhSiXAIEsZP29eiuZq68ZsvYYy1D/06euyxKO4UEzMitgm+Ql5+uc/0XPITGxeMh531auNxet248DRL3Fk+/xSLWLPnP1P9KMvxBKJucAUTERJJTxEK6IXWcaLllMADlewtBdAhaKVkUbjjKGbK3N5u21cdLDN85iVnYuHOw/DkL6d8VyXR3H85A8YPmkJDm+dgxpVK9sVad7KHfjg0AkM7NURbVs+gNy8PMRVruht8QN2/BMPTysM6CVyPYJ8p0yeU979/PCJyYrAlIlZEdskPyHpi+q3v85h7dxXhHDx15PxyNMjsGvNVNxW90a7Al9LSMYzA6di1IBuwkVh7uSXbTRm0kyS0nNlf4dpb4F0dBqflCP8m+UYgMPZp0ijUaFShB5XC1+OyJ2KfJ0peYjwO1VQoXvmqDDrONrSic3pj9YgKEgnVtOu51hB0s91ecxmdfTctnhqBGaM7YfObZspaPXyFfXTh6ZDBVVRkBEUHl+bj6298/nhzyfKFxQLyZiYFbFN8hOSjgGjI8MwYfjzZuHuaNG7RC2YtJTew99AswfuFppKo7YD7RJzbp4B11N9m9PXU+iG6rUI0qmQnJ5nSkUowwAczq7VMoBKTIRexMLOVMjVQvE1VooMQpC26A55x77j2LD9Qxx8pygK1NAJi3DTDVUxemA3G4iOnvgewyYtxrMdW+KPvy9Ar9ehQ+um6ND6QWfh5HbFEDjedKZfMGlxcoJf5nV1UiZmVxHj9gKBAa/Mxa11alp9kRHZThnTG0+0amyFEmnBRORUSEum+0p7xEya5jU/5PT15JaKO+XsfESG6mQZgMOVtVatGILk9FyEh2qdDj7iyvi+aqtVq1Apqsh9au27B3Dok6/F6Y5U6PkMCw0Rz2/xsmX3Ecxa/I54oby19g048/d/WPr2+3hr0kCbZ91Xa1L6PMcaz/KLVfYjX76mCOiYmBWxTfITkr7IyOBr/LCeDjVm6ZibEpFXCAkW7TfuPIwWTe8RWkebFo3EkVaewYgcNxJDyAkd8lkOCdaCNP/cPGUYSZWEHwXp0Os0yMrNh8GgrCPs4muirGAUDIauGFzVmImYt+85hr0bi6JGjZu1GtnZuVg4bYicHj/FyHL0gTf9Imurr8b5ZV5XJ2VidhUxbi8QoDvmM2fPY/WcMeJzaXfMZOj1znsfWyG3aO17aP9YE7R/tIk43hbEnF+gyLtly4XR/WxIkFbEnab1KLkE6dQIImLOzncr+5Sc1k7ELEUDk+6Yf/h4LXQ6rRCzTfdX0Ktra7t3zOb2R9ZBV3gkTi+mWdk5WDZrhJyWqRhZPm40W2R6pL974QolpXq0+enZ+ke/HqsIjJiYFbFN8hOyyCp7Au66rTYWrd2Fg0dPma2yv/nhd8xethXzXh+EG2tUsVmA5VE2HXWnZuaJQB2JafKO9OVoJygAB+U3jgrT4XKisizLLdcmuYAl077otSICm1KL5EuenpUHykGdnZOLRm1fwtjB3dHDjlV2Wnom+oycjX7d2wkL7NT0TLTqOgovdG2Dl1/oiJ/P/IMeg6YL+4oenVspFRa/yv1Rw6L0i9YpHU0JHy2yM3v0c+tvrdNN+hWEUiZnYpbrzshcLnrTXbr+fazctFdISn6dq+eMRoM764rPn5w8jSHjF2H3uum4tc4NpRIzHZNeTc4Wx6ZhIRTAQqHGXxYBOIgAoDIKklZiIUtm6YWptAAqSlibpS855Wym64ZjX5wGGXxJZeKI59G9k4lkU1Iz0LTDYFj+7stvf8GwSUvMfs9EyGOH9HApMIkSsPKVjIfunSPcC61UZVO+KdvfS+08UP/496YTPrkXJma575DM5SPtIzEpFXGxMW5l9KEv/8S0XJEUgYo9H1qZQyDEsxeAg0jgekqO4o6BScOkfbha6B5VWgAVue9NcV9yCqFKLxq0Jkr5d+VaImJjosxH2qWtJ99gwNVrScIbgV5EubiPwMEG86xJ2YqkC1XmEs+33a9v9wMTs/u7xj3LBQKkdRMhW2rIpM2QhiP50CoFCArAQYVSVEqFwlcG6zVIKhZFS+5rIvzJNSrTIsSovQAqcl9HSb7k5JNNLx8imAUXvyCwv/58cz5mKdymlI/Zm5+f+HGUX9br6qSsMbuKGLf3GALkHhWfnI2CYjZSEaFaFBhVoDtBJRQpAAdlXKITAMtCJJCeZRDGYEoo9DJRIUQrNP3ixTKAihLWUpIvuTgBsHCfUsJaAk3GfXcvMF0mS6V4mE0vfX7yfyMVASUTsyK2KZN05ykAACAASURBVPCEJFImrSy1hDtY8qG9kpQl29SC1uRbcgAOIu2K4TpF+AE7ilYWotciWKeSfcYv2hvCvRJFK7PzskT1kaFa4damZq3ZL18ue+5Y6Bc/5g4/D/fLel2dlInZVcS4vUcQcJQ9iu4GgzQqq6Nhj0zs4UHIDYc0fApXWVKJrKBDvsEoAo/IuZCGSe5epR290zF3akae2SZAruuhaGV0SpGeVTLmVR1kn5Lr2gJBrg9uXwxjYRhOyQ7bF587/8rEHAjPD6/BCwjQcW9KZp7dNImW01E2I7qzlas/MClblCaRomM5yk9cLSYElxKyvICmZ4aku33Cm47jDcWO4y1nIF9guk+ndnItZOBFL0MkY2mhvUP0GhGhjY62ufgWgd31Fgvra8mPuSRHZk/Xd/ltqG8X6uZsrDG7CRx3cx8BShtYmoYpjUwBLiJCdLL1oQ0L0QojIrIqd1SoLR2bkr+2HEt0eJB4ASpNw5TkJmLONZiuIuRY6P6YsmA5elki2aktHXtz8S0C79261OqKufjsZcjsKIYqqf/TZ5QRqY2J2bfPY7mfjbTlhLQc5OU7F+JRrj60kguRIw3TcsNJu5ZjABUpAAcZ4jmTPIpOCuKiQ3A5UX4nABWCtaDrhQQnA6LQCUBMuHX2qXL/R+oDAHbUXWbjxlzkMSVp0iV5VLlf3/WPwT5YXdmnYGIuO4Y8gpMI0LEUaTHOaJjSkHL1obUMwOHk8qEP0iBMr0GCExq2s2N6op1lAA5nx5NjABXJ2vp6Si7oVMbZQj7bdPzN7lPOIlb2dttvXl4U4ksK9VVSyC8P1j9zdlDZhffBCEzMPgCZpzAhQMR8Jck5rcwSM7n50BYPwOHK/sotgAoZ2YUEqd2Ktia3ACp0xE7f4XSM7UoxnQAEMzG7AloZ226tvcKUj1nKw+yjn8/+PbCMkvumOxOzb3Au97OQe1RGVj7SSrGSLQ0kOfnQ2gvA4ewGyymAiuQeRQFeXNEwpbXSkXGoXuPSCYizOLnaTvIlJ9sFegF0tdDLH/lvs/uUq8i51/7dWqvMfsxCIba4FPbm5x7/vuSewD7uxcTsY8DL63REzFfKkNSBLGiDdRokpTs2tPImxqUF4HB2XrkEUCkpAIez66B2cgmgQnJk5RaUySCNXv6YmF3ZfffbvlNzdZGFlkTKPvjZ89wA94X2YU8mZh+CXV6nIoOv5IxcZOc6f+9nD6tKEXqkZuX5Lc+xowAcruyvvwOoOArA4exa5BBAhV6Wwh34kjuzHjoBoONwdp9yBq2ytdl0w5qyDeBm717/9Xezp2+7MTH7Fu9yORu54XjC79XfPrTOBOBwdoPJelirUbl8H+rs+I7aUQCO7FyDR4Ke+DOASpEvOblHlT3sKflyS3mbHWHI9e4jsKH6Wr9E/nrhQj/3hfZhTyZmH4JdHqcibZly+VLkK08Uf/nQ0t0w+bzSHWZpAThcWaO/AqhIATic8SV3dj3+CqBC/uFEpJ5KFEIvS3Qyw1qzszvvXrv11d626OjNDMw0TdH4fS71dU9gH/diYvYx4OVpOjLCyco1IDndNSvZ0jDylw8tBeCgl4s0DwYI8VcAFVcCcDj7vPrjBMAdX3Jn1hMdphO+0Ow+5Qxa7rVZF/e2XzTmvpf7uCewj3sxMfsY8PI0naN42O5iERaig1plLDEBhrvjltRPr9MgKowSUbju6uVIFl8HUHE1AIcj+S3rifAT0/LcsvB2ZR6prTu+5M7Ow3G0nUXKvXZrqmyw0GTtpZkqbgkmzVOShZhz9f2vMjG7t2PcKyAQICvs9Mx8pHspbKMvfWiJcNIy84X27+niywAq7gbgcHbN9AITFkJRt7xvOV8WX3Jn1hMWrEVYKLtPOYOVO21WxW4w52M2+UoVGmmTq5sXPw+I7+2OuD7vwxqzzyEvHxPSPezVpGyvLdZXPrQUgIN8de3lJ/bU4nwVQMXdAByurNNXAVTK4kvu7Hoo6AjfNTuLlmvtVlTeWKgxF9d0vfv55WtMzK7tFLcOGATI4Iv8jZ1JIlCWRXvbh9aT7lGO1untACplDcDhSH6p3hcBVDzhS+7MeshIjq4amJydQcu1NktjNhYFFVEVxcSW7LTMinOh3ZanPg++/oJrgvqpNWvMfgI+kKfNzTPgug+OM73tQ+uJABzO7rO3A6h4IgCHs2vxZgAVX74s0XorRQQhSKdxdunczkkElsRs8h0bW7D90EQmZie3iJsFEgJ0t3zNgy5FjrDxlg+tpwJwOJLfst5bAVQ8FYDDlbV4K4CKJ33JnVmPyFMdpeeIYM6A5UKbhdGbCrNLGQuts01Xy5JmXPTTs/XDE3u5IKX/mrLG7D/sA25mIuWsHIPPg2Z4w4eWAnBQwApn8hN7aiO9EUDF0wE4nF0r3c0HaVRIdjGhRGnjC5KM1ItgNZ7yJXdmPXQ3Tyca7D7lDFrOtVkQvdkqpjlhS14c5sQWhZ+l0TxVPyqZidm5HeJWAYOAt9yjHAHkaR9aKQAHEYAb+RAciVtqvacDqHg6AIcri/N0ABXyJacocr58WZLWy+5Truy847ZzIzeXko+5pDzMZf/96OTnHQsngxasMctgEwJBBDL4Ss3MQ2aO512KnMHHkz603gjA4cwaqI0nA6j4S8OU1urJACre9CV3Zm/IMl/YHKhNrj1cyobAWxHvFPNj9kEGCxjxaioTc9l2jnsrCgGDoQBXk3P8JrOnfGi9GYDDWXA8FUCFNEyDwfTC5K/iqQAq3vQldxabKlHB0GiYmJ3Fq7R2b4a/Y9aYpXamO+ai8JnSnbMn68em9vSE+F4fgzVmr0Mc+BOQtpyYlovc/LJljyorUmX1ofV2AA5X1lfWACoUgIOI2d38xK7IWlpbSqMYG60vU8pPuq8OCVL7JHBJaWuRgpqw1lz2p2NW2Ba/aMzj05mYy757PILsEaA3XPJXJmL2dymrD60UgKNW1ELTZRbFJlIZUSFEhy/+KkqwfmuNRVCZ3+yNiIgKxde/mPK8Hv34LMa88KHoajSajudiq4fho29MgQ3WrvoJ86Ydh0olhSFUoW69WOw58qyof2vyl9i27rToJwxhjEY0uL8q1n3QRdRP7n8cx/f/KeQSxahCi/Z1MW1NC/Hx1U4f45evLkBFc6uMYpSO/etjwLQGon5680M4/3uq6E61JGj36feg5Uu3iPq5DQ4j9WJmkeaiMqLnlqa49bE4k/y3HEZGSpZp7sLs9sOvdDJv/cF7diM3izR06WhSg+f/6y7wTMvKx7ct16KA0n8K+VXC2rnhZ0Wp+P57Zi7yDIai3io17j8ySZByvqEAWUPHIF9SqoyAVgWELJlrnl/zZl/T2gr3R6VRIX/sOnN99KaOICNFk3wAdDqk9Nhlro89+CiMBmlpRhi1FXCt7V5zfVzW11DF3ApViAkPLu4hMD10ixTgy/QYSwG/ivstFw7vqfqJGc+5J7CPe7HG7GPAA206+pKLT8pGgWeSR5UZHnd9aKUAHNXD5sFoo/gbMeWtR9HxuXqoV2ORHRmN+P3CCPH7BtWXF4toZPp4+tJgUX9bjcUSJZjHIeh+uzBMfL5P9C9ejPjuoqn/w1XXWH2hSV9qn102kVv7uPXmkIaWE+y/Ynox6F9lh+38RmBNfDdRP6nybhvrYyL5qddM5Luoygd26oGhVzvi3KY/8P3M0+KFwrKoVRoMuPocvn79U1zZ9aOZE6U2Gq0G933SD1cmbEPWmXM2q9dVCEKNjSORPfVN5F+/blMfFBODoCmvQbPwZaiybK9TtNVuRvYL4xH5zlNQFeRb9Rc8UK8Fku8fiSoHHrOyFBYNjUDGHYOQUbsL4j5qgYLgWGjarIMqNNbOPvGvnEVgWui7ftGYJ2cyMTu7R9xOoQgQKWdm5/ssmYSzMLnjQysF4KhbcRHs2fcEV9Dgyz8G4pbq9utvvDECh7/og3uqLbXrVtOq9U2Yu/4J3FZ9sQ2x0rp6Dbwf4yY2xn3VltsQF2kLe0+9gBo1K5RKzBfOZ2Dg/Ttt+pM6sv9KH2yf8SOOLjljj/exOr4bVjz+KS5/n2j7WmAEpl/rjNW3HUB2kjWxicYFwLD4jthz53sw2smJTC86Pc53x/HGq6EqsDUOJA230ecD8O/TbxVq8cVFMKLWrnHIHDIaBYVxlK1bGBG2ZB60s/rarr2QXPPHv42ojU8CKrWd9RmR8sJexO5vVaRJW7QqUGtxvd0hVD7UQuytuv4gIO4eqGPudvaR5HbFEHg9+F2/WGVPyeqhiL1gjVkR2yRPIf3lHuUIDVd9aC0DcNSvtoyJuRjA9GJQVmLu/Gc3fNViHVQG5RMzwaPpegwqXaijR5HrS0BgcvBWk9+y5L/so5/Tc5iY+aEMYATI4CslI88rGZc8AZuzPrTFA3Dce8MyxR5l0x1729i3ZXmU3eH3p3Bh0de4tPMHl4+yNaFBqLlJJkfZwm4AUNVqB3WNh6G6wXS3z8U1BCYEbfWLxjwjp7trgvqpNWvMfgJe6dNSoAcKwCHX4qwPrb0AHPdUXyYr46/7GlfHjsNdxYtQacZfZIU9oPVB/HzyP1kZf3U687T5Mfmu1VoYclwz/rpx5yvm/v42/qLjbGIUOn7XtN0AVcV6cv0TkLVcr+m2+eWO+Y08JmZZPxgsnPsIkLackJqDPINMLL5KWIojH1p/B+BwZQccBVDxdwAOZ9fiTAAVOfiSO7MenVaFmHA9Bx1xBqxibcZqt5Xgx1zU0L4fc9nq38wzeT/IvbDGLPcdKoN8aemZyDcYEB0Z7nAUanctIQUVo8KhD9KV2D713wSk/BQPzYM3mdvYc5Gx1JLerr0H+bkGk0uEyUkG/S8WudhMjdsr5BQmsMKmV4VphVbANMmAqjtALwPkGiRUFZUGa68UaWG9qmw1GQ0VulqoocLGq92FKw750DbUr4YRBeZYvJogNfae7yX8fCkAxy3Ri6EyFgjXIho/JDgYJ872M6/v/morTK5HheOHRejxyW+menKPGiXco6jeZI9cKS4UR7/vK+rfWfUTFk373DSW0eS+VPuWKOw4brrrWjz5S+xY81OhFmZyYLqrQVWsONhB1M9/5WsceOcnGC2+pe56oCoWffCEqJ/z/An8ePg/q/16oGNtDF3TWPxuYvNDuPJ7qsDGNLsKHcfeiydH3yzq7blHdVv+IO56uoqoX3/Xx8iIzxB7I2QA0GlzY9xQ6D514N69yEnPEnLT/lCTLn+YLLypnHpkCwxZ2ebjddqTx04PgFpFgU/y8VfnRTDkF947qwANVHjoxGsiOxm5RyW/NAkFBoNJdhWgNqoRtXZm0XpfHwRDQZF7lVqlhmraCnN9hRWFsgh3HCNUai3SXtpqrq+063Hzc0d3nhqdHvGdLNyjTjVFASn4ha00uhBcbnTU3L9q3jEgqhFUwSVbabvyd0gDZ2ZlIy/PgMiIClb7GkgfXtFs94vGPMfAxBxIz5Gi1kJ/2GNnrMKxL8gfFrj79jpYMmMYKlWMtLuONVv2Y+GaIl/ONi0a4fVRve1+MeRn5OLwI2tgyDSg4af98feWc/hx2im7LjJ0r3jitR/w26a/be49VWo1+p/viDWdPsX5k7aWwLpQNSb/2xGT7j+EK+dSbeSOqRGGN79rh4E370ZGmu2Ret37qmDywZbocuMW5OTaRr5q9+zNmL7hUdwQtggoZilMBDzi9YfQ86W70ajaKqhU1v5TRELfXHxZyFS/+lI7mBrx48Wh4vcNqy23495kxLeF7lMPVV1tx0rbiBOXTH7TLauutWPlbcTRSy+K+p5VyO2kWDQqoxHvxJuI/8Uq9AVoW782/hlRb889ighWejFaEve+zd5S/ZCrHUX/9+put3WfUgGdz3TD1R1/4M/5J2zqNSoVOv72Mn6dtB8Jx3+ywU+j1aHO7qHIenMVss7acZ8KCUH44klQLZoEw/VrNv3VlasAw6YifE0vGPOzbeur3InULpMRs7stVMX9/IxG5NV5Ain3DkOVU01NL4NWxYj0W15DRsyTiPuxEYyaCKibHoJKF2Ezj6t/h1evJWHGwk049f1vYqx6N9fE+GHP4ba6Nyrq+8cZYUept1s9lRae5aK7tz7PKzA993IvrDHLfYfckG/tuwewc99xbF4yASHBQXh53ALUqlkV0181aXHFy879x3FDtVjUv/1m/HcpHv1GzUa/7k+g9zOkTRSV+OP/4vrpC/hz3bdC+2v02UAcuPc95GXYWtqSiwxZ4q69aS8KJI3IYiwiv5cudsakyu/bd3GBCtPjO+HF2B123Yuo/9r4bng+dmsJ/YHNV7ujXZUN9jxkQE4zRxP74fYYcl+yDbOoC9Lg5D8voWG1FXattCtWCsKhH/vh7mpL7Pa//a7K2HroGTSstsxu/XMv34sRExvDHjHT2rZ99Zxwj7JHzFR/7LJJY+8Z+64dYy8TMR9Zfx7bxp6y6361Lr4b3unzNf44cNHmeZCssDff/xFSzmfa1v//VcbQ+M442OQDZCfYCSxTAHT5qxtONd8IQ6792OmtvuuP/7V9qzAIi/UUKhhRd+9oJLw43u7ek39WxTVvwjjpJdMpR7FC/VXTVyFsZTfTN3yxQlycOWgHTNqynaJW43qXg6jyZVO77lNQ6XG18TFUOd3IdNJxU38g+i6oKz5iNZirf4evTl+J5NR0LJs1Aiq1ClPnbcS1hCSsnD3ajW8BeXcZqd5uPsEpomGLyDEUXKfwhMeT9QsKWGOW95MRwNI93f91kNbb/7n2YpWHj3+NUVOW4+dP1juVum7SW2/j4uVreHvBWCuUsuLT8VHLtabfBQAx77vcu0TilAsxt6q6zo5fMrDxq2cFcZdGzIeW/IFd022toIl4nSHmHY9/ivgfkuwy25AyEvODX/bFnx3nBQQxE0DqZsehCoqxwsrVv8OeQ2bixhpVMHOc6TTk/Q9PYMnbu3Fs54KA+7YapqY7ZinVYyH1ml2mpEhgnq9fZGCNOeAeJqUsqFHbgZgxtp8gZyq//vEvug6YgpP7liEyvPR7q7x8A9p0H4MnWjXB6IFFd4VXPv0H53b/jCtHz4ox6Z5YCUfZnWtuQW6e/aPswQseQqMbVsn6KLtV1bW2GvH/RwAjjVmjVuPZSpSlx9dH2WoMuWq6A3f3KLvxyT5IWHAECZ8o/ChbOuqu3Brq6h2hinnY/DXh6t/hsc+/x9CJi9Gq2b3o3LYZ5izfhr7PtsPT7Zsr5avHaTmHaLZZnGYUxuM0n2547/NS1pid3iNu6EEEyIDlzkf6YPkbI9G8SX0x8tl/L6JD7wk4sn0eqlaxfqsvPvXrc9fj4NGvcGDzm4itFGWuvvbVfzjZd5cwACIaaPRpUexouRp/ScK3r7rJrvGXVN+w+gpFGX/d3aQqNhzuBI1GjfGdj+L0h+ettlHuxl9NTvYxy/tXl0UwSNHCCo2/bt470lwvd+OvyqcbieN2YdR47waooxsWHii5/nd48cp19B8zB7fUvgFffPMzgvU6rF8wDjfXqu7Bbwh5DDVITVdQIqC8xU/pcrn476XPZa9fbmB3KXk8AeVQCnpTp+Ow1s1NXxLOaszLN3yAZRs+wLaVr+OuerXMyBHZU+jN4nYyh5tsR+QdMWi89lEblP/ZcxlfTDqFuMaxeHT1gzb1P2z6Bwcm/4Lbn6iOzstMCRYsy6fLfseh2WfQ5Lk66PDGnTb1O2efwccrfsVjL9+OrmNvtanf9Nq3OLrlX3Qeeyc6DTb5moboNULLpLVMHfQJPj98HsMnN0S7Xrfb9B/W60P88OVFTJrbAo91NFkwW5YBnQ7g11/j8fkfRSRjWf9S6/248FcaDvxt/4ugV/M9SLySgv1netl9Qvs88AEy07Ow/Rf7/Ufde0CwwfzvTNbZxcvYO/YhKEyD6V+1s1s//daDCIvTY+SnFIbStiyufQCVbq6AHh/ZD6Cx9Y49CL+zItpvb2bT+erhv/DrG1+jYsMqqD/Xdvz4/b/i6rpjiHqgLm4Y1wbhoTpTgorCXN4ZH36JpH2HEd68ISKfMV3HWJWP30fWic8Q0uxh4LHOtvWfrAP+9zl0D3RF3gO26w/+fC5y//sKqgYvwVDP9tkNOT0WmtT/If3WYUCc7T106JmXoM74E+n3HrOam0K5hgRrhTcAFVf/Dp95aSqaN70Hg17oCLLkfn3uBpz46id8uX8ZtBqN3X1S6i8HaraWdHXs1d+vLGBiVuozo3i56W7r8Ufux4s9TF/aju6YyRVp3srt2LHvODYuGofbbylyhUo7m4jk0+eQ/O1/uGlyWzHennrrTa+uZqMpI27p1xC3jTHFDl5D7k0wWNxnGzHgUtEfxMSw96wx/v8sQTMynjL/bmi4baKFxWlFx+o9wrdYHe/SS/fWtKJQe11FfdEUVL8r7TmE6LXQalS4I2ipjcnQl8mmE4BV83/Akhknio6HycVJrcEvCYNE/ZThn2Dn5l/Mg9PY9EX8c6IpycSoTgfw1fFLReNT2EG1FieSTIZ3PR/YiT/PJJrr6aUnKEiLE9dM94o97tqJC/+lWPQHdCE6HLnygqjvc9NOpCfkFK3faERIRDA2XTS5jw2M3YWszDyLrTEipnoFzDtjsqJ+JWq3yT1JwscI3NQoGsM+MRHUm9HvweQfZGpABjg3tYhD9w9MBLymmhQYwgQByX97r3p4aPa94vNH9deAnJukR4NOJ+uOuB83975H1H//6ByreNjUv8WXE5GfX4DsXAPO9yoKJlI4AWpuLsoelTqG4mFbHt0bETHnbfN+5L/V1dra12iEdmyRx4F2bRur/jRUXt9D5v6hh8nNrHB8esTVRmS0PiXq9RcXQf3fBgsrddpbFTLu/17Ua4NGQK16CEHq+6FS1YErf4cZmdm4v91A4T3R8iETlr+c+RfdXpqCD9bPQN1aNcwyBsJ/+qspVnbhHbIPf65mjTkQHh9lroHcn3bt/1RYZYeG6DFw7Hwrq2wyBKsWF4MxA02GEBNnrxOGJmT9WfvGquZFV6kcDY1KjSON3xI8fMehYTi78Bec22T6IipeWn7/PL4a/j3+3v+3baVRhe5/P4UNrT7D1Z/sGBSp1Bgb/ySm19uH7ETbO2G1Gph+pQsG13gf+XbujPXBWiw+1wX9quywmwghMiYY715+BncHLbNrpV339opYdrQzGldfacfOF3jkyVswa2VLNBY+zdZ3ukQuY2Y+gq596uHxuLdtDOyo/tAVEzE/RnfGxWYg8jpSaGX9RNX1Nu5JRI4HLps08z5VttmpB9ZfNe3lkMo7bd2zYMTSwuxRr1febSs/jJgWb0oruZTco4pbqRuNGHylM/5Zfx6nZti6xtFx5LN/d8Pvk48h4eN/bfaeXM+bnuqLc6PeQc6f8Tb1aqjR9MRY/P3iFOSl2brGkSYaufINYNpQGPNtrcBV2iBg8hKELu8O2DMCD41AZr9VqLitA4z00lGsGCvWQfLjixD3WTN7RtzIrdoSSXWnovKP9wtr/uIlveZYZEQ/hdDw26FCECJ0/4NKFQRX/w7bdH8FtWrGYfbEgQgN1gsXxk9OnsbejbMCTmN+UfOuFdaSn7jAtvgVs0US0bLWrzNwrGxlsloASE1v32OmrcBnp34Uq7nz1lpYMnO4+c64c9+JgqjnTzFpefSFcOGyrT/o/gmvIOfHy/hv+3eiXcNPRuJQg00wUD7d4l9uRiNane6Fd2/dDVW+bT19eXf/uytmV6bgDXZ8WIzA2OsdMaHSbtj99vv/1Lkz47ugf+wOu+5LpLlShqTelXdArbYdvwAqbE/qjoeiVwktp3gpMKpw5EpfPCDco+y5T+nw2T/9QMFG7NVHVQrFoR97oW2Vt0mFthn/5vqxWHLoCTwaZ88vGWj2+I2YvP5RPBG3wdZFyAg8N+E+tO5YG2Ma7rNDnMCcb9vj96MJ2PbqV3bdo5Ze6yrco/606x5lxLRrXUDuUan/2XOPAoZc7YRd9d9DfrqdvS8Auv/TFV88uMGkbds+HHjwVD+TFXaBffemOw+9iqt9xsJozlNtMYgRiF79hnCPsvvWhAKopq1ByOLu9t2rVEZkDtmG6C1P2DWkMxoLkPzcQVT5zPZYnqQwQI3rD3+KKj+SMaWt/NCGIO2+L6DSUU5rFYLUfRGkaYbsrAdc+jv87c9zWLFpD46e+B6hIcFoWP9Wcax91221A+BbyXoJfdSFp15Wd8wl3S177vfrmZgD7llS3IJS0jKQl5dfYmARRwvKOJ+Ik0+tMTUrAFp98yr2PLgVGZfSbbsWAC1/eB676u9FXrod39ZCYp5Tlfya7RGzCmOvd8AE4ddsz/kUmHm9C16K3WH3u5HoYM3Vbugj/J5t+1NOYXIRekxotLb1mmAtPvznhRKJObKQeIm47fk917urCjYc6mxXYyawnhl8G3pPbGyXmOm7acPXJven9nEb7LpHSfmU+8Sa3EwsC2nk6+NN/plDK9umfaTxiZgTzmdg0X0f2SX+qdc649SM3/Hdst/tEitpzEc6fIzrvyTb413x0vX1E1uRl5hl59kw4sGv+uHPLovspoVUGY2ou380UgZOQIFtMmzh0Rq16k1gMvkt2ymkYU1fhZClz0BltNVpC1RGZA/Zhqh329vndaiQ1GMfYj97UHgb2BR1KK4+dBhVfjD5Ldu81OlvgLrJAaTl1zafZoTpPoVGdYNo6urfIb1Y5+cHduSvFzRbHH39eKV+o4HzMXsFWB7UNwgkfvMv/t30NRJO/WOaUAU0Pj5KfDHtqbfOVnPQqNDym57IOJ+BvS0oRKW1nPrwYHT5oT3+PZGMbZ2PW5ELEUd0rQi89M0jOL7kD3w87X/WX4BGI+o0jkXffc2wcfjX+GLrv1bfj9T/we434YVF92Pmk0dx9qtrNv2fnXIfWg+qi96NtuPSf+nWd9D/bxg2e0c73NusKlrVW4/M4pHEjEbsPNVLEOeDNVbBIGJzFhU6av6qLJtC0wAAIABJREFUMBJYm7i3bTRqy6PsR+PW2YkEBhy5YgoY0j5uve1Rc4ERB672QZBWjR7R22xOBAoKVNiU8IwIW2qXmAtUWHrddAdt7yibrPqmXjcdZS+LoxcjW/crImYq22rTi5F1vVoLdDvTVez9D9222dRrw3R44OPnkfHdZVyYstUm4lNolf9r7yrAqzi66HmWvLgbQQsUdyhSKO7u7lAo0OKUFikUL168OMEdAhR3hwLFHQIE4u7Je/v37j7Je7v5C20IyevM1+8ryeyOnLuZOzP3nntdkWdVL3CHziJy3yHTq3qywRfOC7vRgwD/NdBeM71KJ+zllaoAzfvAesdEKN49FsleVr0z4su3hKv/1+Bi34r+gJLqTUeiZ2n4XGwMbZrpppPaD6ngB9gVgPutqlDITO/K6VNI/vImP6cUWWHDd62U1YK1og+U8upZ8webw3rpptiYjsdsDCaiDyoihGDV85gzr96PKeYc9qWw4ZogEHruKW6NIKcZClAsQ4lDQohJb1c1Lv98ES/W3BdON1RtbYXal4wRdW6MvYqH2wMM8YWtba3Q9q7gfERlV/erePL7O52+56B2V2PoQ6P3629NzuPl9RBd7B/A2dcWY24KjmdUJtU4inePIgXqFsfBp4gLJp1rYKgfXc4f4YHxBhpGscpeGO1vjMrUpsRmxIXTyU6ITV2lfn5M8TN659YrvApx5ECl079N25fE+IXGa87q+VcjNVWopzH0HFASgyYaPc+b5F4HLZ/ggwPF7e476Uu0HUDXnEKpn2s17zSljwv+48qaqN2soKG+mc863cmQFicFDrwVPLfdnayRkKxBJ+fNBlsphTbdEtkZapUCkbqbiiEeO01uBRaFtDe0fetAEPb0putu3kULFB7zp2Bj3PJbK17j4qQbFFmcx1+mkGPQG6Psbs6+h0fLH/D0M1JGKmsV2t43ekY/m34O7/YblaNKbYXKp7oZ+g+ZthuRl4SNFSGgsrfFZ1sGGupTZv+G+CfPdZEwZbB2dYHdjDHGb3Ppz+DevTH8LPPJDQyaaPjZYcNgpEULIV556eYqhqS2PxnqnfZ0giwh1hBpU5unImJrTDbU+1xsBE1qPI8fsRBSvZsiqsgPhnrP29UAbarw3UOOJN+eUH8+AiFRyfzGyN6+pEDNkwH2yrVQyllaSKPwjP/qqtiYsY3Z3KaciT9v1hi/RalxZZffsZCc2UUS2Wgc2r+cfUKjkqEx40fdaLIVDkVdUOm3ZoiKM3XQCj8QhFszrsKzsidKzRMSKKQvFFP76uzbyFszH6otFLy30xdSCOfm3kepTnlQ62cxfeno3Ke4svIBKvcvhga6BAzp39878T4ub32BBiMLo45OCSoVcrg7WfGL5rJvr+DWkUB0mVAeNboKsYcdbVWgucYlpmFO/8u4d/4tBsyuhirNvEXj+7njKTy9E44N943JM9I/RCf14Ocx+PWeBH0HwKj6BxERnIg1t6Xfp2QTiXFazP1Dmt60oOYppCRoMeaaNL1p7pcnecUw+pJ0/fLSR2HnZYXuxzKgP5U9AueCtmi8S9rOevSrA3DKbYPKm6Xbv9HSD3YFPFBkgZheFH/xGYIWn4RdKS94/9CCx10ulyFKt5FIOn0DyfsOw65kESj7toWPqw2CIhOFpCFUzp8GTh8CylYFmonxlV/2h/rWASSVbQZtleYi2TleXw6rN6eQWLoj4j8TbgbSF+dHM2ATfgEx+QYg3lf8vuvrH2AddQsxnw1DvF1Dk3fdHK2hPDMM8riHeFdjn0kdn73M2VrSJyEb/bl/kqF0Uvh9Eq/sLWlMMX8SgbNO/x0C0bdeI/bBa7w79xS+U4ST1qVqa/nTcTqGDQq0K4VcI7/g638vvtXM9seh1PAKyNNfOCWuy7fb1OGLA3q9MtKj5rjvE10/jgozntJ+8CR6len16QydFzG1P8zTSIehn2k9XxjSDrRoEgWntZOfyLZIcbSp5HKzQTUlJZIwUmTorLgzWLBFbZx5DzsWXDdcr9JJl05Ke4K78/WL+l/G1f0v0tVT8isZNrwTvKS/b3gEj/4MTkePApRKOXYGCgvE6EoHERoQm46+xMHaWoUVrwUFPrX0QUS9S0rPTAPdQEwNEBTIpEIHEB+dbKQncYCtmxUmPxTq5+U6gLSUNBN6lXsJJ/Q5I9wgrMm/D1pNmsn1sW9NLzTaKNwAHCi+DbzZVw8PB/g0yYdKus3XtRordadr4bsjfHJ3qwDfgQKH/nHTBeD4JCBG+dW88ANCooUTZmS/8fz7hsJxyLPxF1ipSHGnInUsOSiact9Us5YYHlcv6Gp2Nc4haRgl9hCKy5amJnMjZR/Z5YCh3vPkVyb19BkE1T7D19u9Wwv7oOXGsdENglyGoDJX+N+5PZwHRYC/yfDIMTCovsBvdrs7FCrfWoBzYcjdhSh8rAgIdFD4fRIotmuEv9vsXtiJObtLKIvHp01Ow4X60/mlstCe0Qja+hDPF1LaQjOHIwBfXuqLW2Mu4u0B08hT/JBlcjS51wH7mxxDxD0xBYacbHoGtMbS4keRECrOACSTyTEypBkm5NuPtMQ0EQpypRzT3rbCcO/dkhQYa0c11oa1RWP1ev5UbF58Cjrhl4tN0N5ro6QzWfk6Pvhhc1209F4vzq4EoN2wCug2tgS6885YZq1zHIZvqo1y9b3Rymu9pLPW3mCBl9yTpz+J318XImwcRvObDnObLzA7VNjYjPHYKUnfmh0qKPZZ7nslvbjH6DY+q/OINz2EVr/XbRByLAhXvj0tYXMGmj/siMDl1xG4SYI6x3H44vwAvJm0F/HXdT4KZlMscmA4osbPhyYoTCxbGfD59jkIGvkttAkS34ZcCeX0BVAv7gVZmphaJ1NZI2HwGrhsawazxGB8X3IbF4S18oP3qa8kY3Vz9oUR8sUqeN+qbDy1pxtlilsjROWbDI/DtSUd9eJLDUF8rrbwOFudN2XIW17mueysGBFoq9jwSU7MO9OYYmbfYQ5DIOLiY4Sde4Tggzf5Y2fBvaN5T9u0CDGFhurr3RyAPcW2QSORVpHqmzzohPX5d0kublRPp+bZ7vulKS4cMCqsBca675bkHRO0dGoe6rlT0tOWTnnLozugvfNGSXoUvU+n5nZewpWaeSEb4/Z33dHSc4Nk/yobBXa86CqtmOm05arGsget0Nprg6TiL1DWDfMON0VPzy2S/Tf+thg6ji+DUR67JJ3FfvhDuDKeWeGwZD0p7jNTH+LKwkeiudGp9vuwVtjb4gzCb4aLv1KOQ983bXGsxl4khYlTapJAmz/shOt1VkOrD6eZrhXaB1U+/zWeNFtgeho2PMOhyIERfPYoKXBoY5Br3QzEjxKCukgV1czFUC/oIrlpovElDt8sOi3r29Gfmr1OfCXtZc0BoXXPwOtmZcm+NXIrJFW5BNt95NglRZ9yQkidvfA8W4OvleVvDeSqAbmXMY52DlsaMn24bZQbTKNxZh4j6v+2u0cjHWkv0yf4LxtkJ+Z/CaAlvR73NAg3+/xmmFLBPaNxr9ceRD8S50umxa3d48HYUGIrUsPFpxqqb/qgEzYU3gutVOo/LdDrdVvM9vTnHbjE2kOvmPfw0ZfMCynemWFtMNR9p7Ti5oAFoe14xSlFv5LJgQ3vOqOt10ZJXrS1rQobn3dESy/a2YuH553XAcuvts5QMZdrkhvD11ZHK+8NEidioM/ML9C8d1FJxUyKc971FnDLa4fRHrskeclzdCdmqpdwouZP1ESPWlX+uCQ9ik7Ml6f+iXsrhKQk6QuJo9+bNjjb/hii74plT15RzR91wp0O25H4TkyfIqe4SqSY2y+GNlF8oqVNWZGDwxHWfyJknFREEA5uq6Yj9ftvpalz5DRGinl+V+lNnZZD4ghSzBnRozhEdD4IzxNfSW6KNEoVwr46Ds8bFAdbTJ/iFM6Q1zwDzbYqkhuDNKcSiKi8GF5nqhsUv7zOVsicxKFdLWn9+JC5tFQIN0n0rWfl//emMcX8IXJiz35iBGLuvMKL5ScQc+e1MBIZh4K7BW/YC1XEQTGILtroxkDEv4yDf+09IuWjslOj/rVWCDkXhYNdT5hHUYRjbge0udgAVxY9xtnJD0X0p3xVvdDBvzL2Dr2Ny5ufiuqrdCmEVgtLY2Hzc3h+Och0geU4NJlclqdHkQ03KCDGVLlxHMbtqo2iNXzQo9AOJMaZnwo5zL/ahqdHtc29ERpRwBQOe3UhMnt4U8xfc14x4KfjFbf23CAKmEKKb2+wsEBInpi1wLrQTjw96jtn2niY0bM44O8U87zwdrzz3i+8/V6seCmYC5VVuXdLKDcZ+r4WnKz2F5XgTUOGFg+EEKlXavwmkj0nU6Dy2b48fer1N2S/N+1faWeNgtsGAfefImweRTozFvKktvZ0gcP00ZCfOIKko3SjYnyCkLDKXwgYOAw2R5eCu3/e1IZMduDi1ZHYYBCcTn4PRdA9MxM1h5RyPRBfrCO8LncEEoJMNyXgEFt5NhLtKsPnTj1o02JF9SnldkNrWwBW+76C3Jx3zXEIaXiaf8fzbHXj2NwrQF5sAGTuQrjN/3pprlz/SSDwTxNMSNm9sBNzdpdQFoyPdq3hp+7j/iThWlihUCD/TmNy9reLbuDVppt8DGRa5VQqBb4414sfmZeLGicHnEPA/qeGpPVKOxUaXjM6d53pexXPj70W6FMyDjYu1uh0y+j9urnJeQReDdctYhzs89hj4E2j9++CGkcR/Mi4QHoVccCwdPSoieV+RwzRo3SOX59VdsPQdPSob4ruQSwFvtBxdMo0zI3Rfkbv4675tyE5KVUXCpBDjY5F8N1CwbGNSof8G5GcpNUdvOVoMbg4+kw0LrB9fLYjVaOPPy1D55/Lo0k6elTrXH6ARnBxItv696vJ89tIj+rps5UPJSnoHxnW6cJrEj0qPjEN3zrv4Ck4fC0nxy+hpp7F33vsAafj18o4BRZGt4OttQIRsSl4wtOjroGiW/HvK+QYE2zE/uGKx7gwlWJ/E/dLkH2vAKPjXcCGx7g7/RY0RI8ieppChmZ3jXHL6dt4u+0PaGl8MhkU1gpUOC5wsqmEzTuOsJN39GmB+LjfBXcMMdQnLVqPuFuPeWoWzV3paAen+XTFLRRu2VykvnxuoFdZubgD3xvpTXbrx0Ab8Yb/9qgNuWtuxPf8xfC+277e0MQT9U64dda6FUN0A2PsbZ+LzaBJjOXbJ18EjVtVRJSdZXjf53YtaNIShHqKs+PTDMpi03lPf/q78TzeAFpNioANgKQC7RFT2HgF73m+FqAhHwkZFFVmAT513isnehb82X/SLprqFLP+xKzPXGH8WSd/3Yk6s+oPMsX8SeXOOv8ABLQpqQiJM80edaf1OmijE/nDoEeDIvD9vo5Ji5fqboImNgVymRwFupeDzxBTCpR/6d1IS9JAKwdKjiyDIjoPbX0j6z7bC22qljQBGiyujNxmFKVfcx2ANpXj3x+RTpHo35/mTd7EHOQKGcYFCYqETpguDlY46heATX0v8t7EMpUcv7w1Khp67uSKx9g78S5fr7SSY36gwON1sFXyp991P93CwYX3ecchtVqBJQGminDz6D9xye8JiKpsZ6/C7Gem7f/W/zL+8A/klaG9ixrzdB7S+rEv63gBz0+G8DxbZx9rjL9tmiFqWfNzeHlZiCntnNcRP/xR3wT71bVPIPiusBHxKOqI/udMKVAbvjiJ8BdxvKLwLu+OTkeqmby/o9xhxJL5gTyoG/iiwWrjJoTHp+oepESl8Px134Z5RfS3K/U2QpOYCjnHIVebksg7yvT9wG5zoEkm8LVwa1sDufrU4scSFS9cawf2nAkuJQ0U/9yhYz04tDCl12l+HMbLlq6RlR27wrtOdYMXN71vs7C7IV1g4lCxd6/zlpaQUxIVhQrhHfaYzN0OR2D/aiyvxLUKK4T4XjWtJ0/sVyv5WwotbBBSyTSDlHvML1BEbOPThMrkNniX76LJ+85R42Cd4g9OI4NWaY8wjwuGeorU6uVMfpE2H/DXaZmPNuYVs3ATpGf0Z8XPv6cJBwrzEhYRDTtbG9iorbIF4OzEnC3E8OkGEXX5PuKfvsO7bWeQa50QSOFmnSVmkbk4WNvboPh+IRHDucqUzceUwuJYyA1lNgsKatfn200mRFe3XpV8UGOzcEpdlXeXKLpTsTaFDfzmeSL6FPDlT0VR+VuBfjWF9zRO1wUHjA9rxQfg2D/zHnZNvCW62p6j81LeNvQqzm9+ZXo1TrSnEMGLeXabU3hwPlhE31qju5qe/MVhBL8UTliGwnFYoksSMezzvUggpWYCD4eVIQJ96oc8e5GiS2+of5/Chc7RnYLHeu8ShZsmm+0snU15qudekcmcNk/jQ4TNxXyKRW4WclQuV2BokJA+cU0eSmJhip1KLUf3J8L7h4rp8uTqHuGvjh3VqH9ZqL9Q1TRyGZ1wHPK4ovR2YfPyot0Mkd3VsWxeuE/sytOj3nSeKrpad6pbFg59hfGljKFANqZ3305Nm0PdtCkfQMVmXjeRWSJxhBDe0e7BNqhuEjXO9NuM7CrQo5wjJ8MqZreoPji/EFPe/cFgKKJvmn/aCKkiKF/vVy2B1ACT8dHmKrjgLeH9kOqQa8yvvoFQH6F9l6QqUNkNgNyuIuQq081MOon8J/7ZQCkkazGJ9JUFPx8xU8yvAoP5JD8Bb4J53Ns0+QoTR/SESimdZnPWki3YsOOIiYzKlSyMjYvHZarcmGLOVDhzXmNpMQm42X4qv2/13TwR0Ycf4sVs01MCzYoW4HInB/MZhEKP0OJkVjgONa72wZH6xxEXkIGz2JMO2FhqP5KjxfQn2jb3fd0Gv/oeEni3oiLHiNBmmO7lD61W7DBEymdBcjv0VG3V8WZNG1Cp5JgR2Brf8lGxxK3bOVphxtMW6O25TbLet6ADplxsgsF8rG5xAxVbF0Tv5eXxdQaxvLvPq8wHNhnJc7LF2M3VbRzGSDh70cn2F139NPe9khSdcWF6xSxtUx4RKmya1hA9SszvQp9XbfFuTwBujrskHh44NL3fGS8mnsS742L6Ex8S9VJfBA7/DamvJWQPDgV2/oDAr+eAixM7Cso5GXJtGYe0H0ZCq5GgP5FSXL4C8VM6QiaVHUouR9LQDXDe0kLSkVCmVCCiw154vy4jmURDLnfEuzxn4XGlmnT+FLu8CCm1FZ5Py0kyAJIdGiDScxY83pWRrI93HI94u/ZwiSnFbxit896ATO6Q8xaLTBxxXaUudaxg4DLG9uLNTel+zuT6E2nC4UJfvh49B/Z2Npg2tj+CQsLRYcBkTBzeA80bmN4y6Z+fuXgzXr8NwZhBxjS2FHfA28M1E9GhvbUQG5CV/yACkZceIvz0TUSevsP/bfhumYA7rVZDYx4rmlfMQLmTg3C2yroMvKg5NLr9NfzymqZz08PKe3E/6YDVufcAUl7W5An8ug3mulP2JGkv7RFhLcSnZV0H1P7U6LYY4kxJLCSE+Ve4xDlh7fAdn+RB/ABday8Ka4c+vBe3RAOcDGtCOmCwpzhJBPVGV+ILX7fNUDE7uFph7oNWGCmleP9SmHnLu2Po4Zq8F7ZU93W+K4RKPQphWfljkl7Y39yoj7sbXuPGr2J6FMl2eGhLnh4VcUuKHgX0ed2Gp0dJediTF3aTh51xueZ6aFPFmyZaQr681A8B7WdI8oIJn/w7xyKw81RzPzm99JB78wSkjvlOMkkFee3bzl8C5S9dJNlJJPukERt12aMkvKjBIarLAXi9LM3z60X7IoqHne9PeF2RXoy1f6UkTax5GfZ3S0rToxQqBOW/Co93paXpW3J3hHidgGtMKf7bktu1gNKuEeS2piaK/9ISVEtJNy9Z75V9KtWomCm5SLXmg/nTLp16qUxb6IegkAg+G59UIcUcFROHmT9+/VHFxRTzR4U3ezceey8AD0esMCyO//bE3PHZYKwt7//JTsw/BjfHaN4RSpx6kJ2Yc/aJWfXLr7Bd2EM6n/JHPjHL7PICNfaCu14sU07M9Adn5bMTcusy2XuB+Iij+0pFecmFoj8fZ8XPZ1P7Gfp99jIQLXqNw+ldC+Dh5sz/3m/nUew7cgE7VxodDNPDQIr56JlrqFK+OFycHFCnenlUKG2Mg59ZkDHFnFlI5rB2Ep69xYu5u5Dw7B1/TS23Vv0rG7N9IRfU2NOejwu8NrdpSressDH/FNGapwdRdir/n2//YxtzRtmpssLGbK2SY7jbjky3MSvkCowIbYE0jfaT2JhtSvjCa7IQcemf2JitGjYF6jbklaL6E9iYU+tfh0bDweZhI8hSiU6Ynr71YTZmWSo57ckgsyoFldtkyK1L5bCVI3OGW01FNLustzFfSO1vmMDNu0/Qbcg0XPRfAicHO/732/1PY/mGfTi5Y77kRP2PXsTLN0GwtlLh7qMXfO7seZMGoWGtzPUZYIo5c76zHNUKKeKIM3fxYsZWyIm+VK4wnEYZs0PRZD7EKztXh9LIP7Icj4GnsxoRscnYU2JX1nhlK+WYFd0WMQlG2yRlUNrc99IHe2Xrhbhr5r1P4pXt6WyN6PhULGh0JlO9snuerAF7GwXCY4Q82VntlW3X0Zh5i/r/UK9slBPiblOxt1ECszoLCVZkMnxsr+yoqqd5T389PcoxbBZs4rb/I69sGr9r6iAg5RzI4U/pvoS/zpaKPJejFpR/MNgqKgpklNFZ2dzGrO8gI9vz+9dfTh1gGK3+xHxm90JDzvq/OzGbT3Xs9N8QFR2L5bOM9NJ/AIfoFaaYMwPFHNYGl6ZBUHSKiT3veffl0EYl8LF9basV5LMApS+3m6+CJj6Nr3dpUgx5R5mGFzxbbQuIdkXm48KDv4BPn6Im7+8vvpOnc8oUHMoNK4N86Xi+9OCqAvvAkd5QaFFraVUUNqNPLfQ+CI1Gw/NshwYJ9CI9PYoWzV8894OOmgqVAiPeCh6++kKKenvfK5BpNZBZyTFVR4/S19Mp+9CUu/zNgdJWhRkBwtwdbFW8cd1v9G1cW/6It7NbOSgw6bkpNvuHXMeN7a95+6qNsxXGPjHNELWp+1U8OBzIc20dvdT4/q5pBia/5ufw+nIkLw+XvDYYZEaP2trwIoJuhfHnNJd89uhx1ZS6trPOUUQ9FsKmuhR3QNujphmg9lU9iISQZJ7e5FHNG3V0CSr0879UdztSopPAyeTIVTc3Pp9pSr+602wd0hISefqYR2Mxde5Ft4VAvOCJ7tK0IvJ/14BPmajfLEUPmQSkCt+bU4tGkDU1/Xa48UOhTdPwtnN5h15Ir4hpjKr5fSCnIC+cHMlj1hg2f2l8as2/KGVrOoLj0iCTKRHVZ5uJ7O1CzsL+Otm2yWFQgZDGR03r326GfcBSgT4ls0JIlVMm9e4py6AIWMmf7uQKNd5+ZqQ/0YNO8dOgjt7Bfzuc0gGhXhRX3lhck0dCkXxMcDpTOiHM/qyhkqezuSogk2UPik5WLmOVVMv5jZWBt0x/XFnw87UUo2KWsjFPmb8BIWGRGdqYzTFasHIn/rj9GH6LjNz7zMCRKebMQDEHtRF9/T7i7z5HxIGLcF0xgR/54xZzTW1nfzlSKR3VKLCJMvsAf9SmhSs9BQWwye2CYhuFU/ap8maJGjjAtawvSq8RFMTuwuTJbASJltMCbfKj/Ezh+mdFbtPIYfTHWnVsSZTW0aPmeZhnnwJ/NasPwDHJ0Sx6FQfokzSQ0j3y8x2zUwmH6SFCAJRtQ2/j2uYnZg5VHGbr6FPLG5xFwM1QM4cwDjN07/9S6hDigyj7k6nD2KQQIXLWpPwHkRJvzP5EvyPftqk6etSMXPt4vrZpkWNcqKD8F3ofACcKWynDsGChfnXevYbgIfo2KEdz3wDBS3vLZ0SPMrW5K1RKdHgo1J+qYC47DionW3x5Usgs9ketZWZR1zhYuzug5E7havpZawrmYXq161QoNzzndOUX3civx5pR6wCHWhWh7CrgnzJ2sJnDFAdV/Rb81TUV1aweZtQ9QDFhoyGAiuPKtiJqWmx/ymYG2L3cBtv7K0TfbkjTE3y988sZsA48IHIoC64mKF+vwK8hi7tuJhsOQUVuCvWhzcElB5h8G6RfQnPr6Fdx1cGlRpvMn4KYRLpSwBXANaUS5Da9oFRXhFz938rbXN5qmRmuWfPjjZRvTDrqN2o2HO3tMG1sP5FXdmxcAnoPn4W+nZugcR0hbvr833agRYNqyJvbG4+evULvYbPQr0tTDOguThf6b2bEFPO/QS8HvpsSGoUH/QR6lMca+j/wpMVciZlwKLx/FF5MPYnIE1KevhzKnx6Eq633IeEVLT6mhfiJtf/oid8r7kWi7grV9Ak52jxuhzVFDiBNF3gifT1FyOr/piXmeR/gT8LmhRThtKR2GGO9Q0evMH1CaSXDiMAW+NHDjLere0ztqMTEpy0wxms3r0DMi3chJ4y8WA9jPYn3Ki6l2uRHl+XlMcnTNICF/slm88ujYtd8GCeRspL6m67jJf8d/WmBl2mOX759DhgWIjhzrZQIqUmz6f9a4BVvKbhDnBmMA7o8F3jbpytukJR9res9eercs1mmJ0hd96hw6htEzD2MyPOCkjGXfsE9Y5Awfj6SQ4RAKekLBdpwWjEDmDQGqUniBCnEIJXzsbB7C0FozAp57dv9vBVp8wkDcT0nlyG2z054HKqniyhmNjqlGqENDsL7UjVJL3KZ2h3BFfbD60E5SS9wjX0ZhPquhWcgBdURfx1Jjj0R4zACrhECPcq8JLlMQDzawzmhlGA3z3URMoWnBI6W+6syVrTZNyFKGSarZ0x9jPo/U0wTo7x49Y7nMb95F8r336pRdUwa2QsqlRLRMfGo1mIwxg/rjs6thENGxwGTeduyvtDzE4b3gNo6c289mGK23G9fNLOYGw8RdugCYq894P8iPNZOwfOBv0Hz1jQogv7FwvtH4laDFZLdrtP/AAAgAElEQVSLI+myCqe/wZkKfnyQAIm1Gc0e9MF6380ZJplo86QDVvruBpcBfWnA29aY575fFGua+iLl9kNka0x13pMBhUbIoPSDxx5J+hXpekqCIZW9iZ8Lx4FSJ471kIolTVfywPR3bTDJQ7p/K3sFfnzeAuMy2Bi45rfHyKsNMNU8WIoOyCLNcqH25BLYUFGchIKwrzahGBzt1bgw4bpkLOwuFxvi+rSbeHvcNBa0MDeg8/N2uNJ6DxJfi2XPb6qu98TNZqvB0RW1eeGETdnztnPA0Z21RKEEKJH9x0pT0+i0+NsMpI41huY0aYLjQDmXrWb1lJYtfXE/+MFueRtJ+hWNP7bfbngerCv5PgVHDWtyEl4XiR4lQZ3jZEiqfwU2N0vwIVRFRa5AUOFr8AyU9qpOk9kjItcFuEWSY5e4fS3nhEjXc3BJEOhXcpv6UNl1htzmv5N9qqSVMae2Cb56bZzRuvwv6++mCLeA5iU4NJLnM9vZqv9WI9BJOjI6Fh5uLh8tUhhTzH8rBst5IPbWYzz/SZ896uOfmLu8HIg1BbZ8tBPz8ODmgm1Z4sRrySdmmu7wLDgxv115AUGbb0vsuThUODXok56Yk0dvgOMqum7P/BOzXO0Bru5hcJeKS3LKM/vETABbefhBoa5qOYvN38ykmNWiT8Jjvp+cwWYwmyHPFHM2E8jHGk7y21A8n7IaKW+FxPQKldVHtTE7l/HFl1saQqmQYW2uTZluYx4V3pIP8Uhllvs+UYjOzLIxL657FoG3P46NWW2lwCT33dBksY1ZqVKi69M2SEnTfhQbs12BXPAhWhPwUWzMqd8bMxN9DBtzasNrPDbqJ72gjL9h9ieZuTZmLk2IhCZTFoKV+2LIVf+N1JBFrBfpwnEaY2UL4Tk/7s+Pk7/7WEtsprbLFHOmwpk9G6Nr36hTf+DNgq18YgJVkSJwHNPVZLCZ6ZXt278CCg4swbfv4WTNJy+g+NmZ4ZWtVCrwY1RLxCWaRqCa7eWf6V7ZeoD2Tryf6V7ZdHtP1LKouBSsanQmS72yG22tAUc7FcKihXSXme2V7dy3tsm3ldle2ekbt1MroVraNtO8sqNrnIGzPdGjkviLGIfYxbANWvNRvLJpHo6avlCmXuW9tlXuc6FQN/9P0KcKWi80KmG9N7ZeKeuichqUdCbWP02WjuiV3VZuppizm0Q+0njehicaWo6/eB+Rv+6DViukIlQqFPDdQt6zQondfhsvV58VdrQUfESpRIlDRqcJ4YqTsjPp7IsqlUmqv2eTziHwwFMD/UFlY4Vq54wbgcuDziHo5FshwACFW3S1RYNLRq/GI63PI+hmkD7ZDOzz2KHt+cagABy0aC4oeRiRz6IM43Up6IzO6TIsrSl3ApGv4wzxd32+cEOXQ9UNz/9W9DASI5J04wMKNfRBMz9jgIA5+Q8iLS7VkMaybNcCaLywNM+hpQAqE133ITUplXcsogxIX3xbELUnChsRKot8DvDULlK+tLh89XN5lB2Qx1ify59PA0mFbJgNV1cyoYety7ubskDy3tvURu9XptmtdhSieNe6K1xOjvZPjSk2qc2jpbYJntrUACdHw7tGjnr4iWd4NOGM4PAm4/j5VL1gTNMYu/823q04KVwQyzjI5EoU2U+JJYQS73cSMUfO87cVNDe5tRWK7fwJRF2KTUhF2qZdSLhwiR84fTsKtRq28wQnQ76s/w2aJzd4Hi8JWGZnD9mPxjSM6i2ToX37xOhJ7eyO5L7GYA8u+4dAEx5oaE7h5gtlj5WIik1FqkYL79PdoIl7KwBPn6drSYRW+9XwvM/VZtCkROjSPAIatxqIKJYuzeOL6tDqgoDQ+8meLRHl8pPhfc+gCtCmphpkm+TUi3f00hfP8ArQptGmkYNcLke83RjE23YxfqspZcClUgpNQbgpNnMQD8EDnYqPqxVk5MBg4aWA9QJhhuY244/884vkYTkCWaaYc4SY/vkg46/eRNydx4g+chb2y4QF8FXHaWKnFI5D3u1ChpTbDcj+Y9onLeSljwrXQEShMXdqIU/YiicHIP5VPK632SF639pBjconOyHkXBQu9jtiFpkLcMhrhzrHmuL2ose4OZfsmun5VRw8K3qh54l62Nn3Ou5tFpS+oXAcinYoiNoLS2N78yt4dSnYvBo1dNmp1lY6heiAGBEFqNWuWshfwxkLCh1BUrRZogWOQ/8b9eGW1w6zffZDk2bm8PRXrOaxYYKX9K/e+8Ve3n9RZIbqsj8t8twndmbjgG+DhfdXk5e1ub+RVoY+bwT61Y5CErG6OaD9U4HedKTEZv5WJH0hHd3onhB0/0K11SJ3JNLxVS8Lyvlh0/mSSS6KHhjO1wd2nSyKB80pFMjtNw5ccAhiJs0WyV5lo4Z6zhTYvLiNuNXLzKh5HORu7sCIqcCNfbA+JeQE1xd+A5GvGJLbj4PN1aVQ3T9mnhwLqSUawrbuUMhODIY8+E9R9qz4kl8jPn9HeF3vCCS9EWWPii03D4l2leHzsga4FCGdprF/ILgIJR4pAPfA8pCLqGscQnILdni30FKQc+bC4xDqKdQ7JZTk4wCYFBkQaSN4truklITCtheUVkSjMuW6m76U83/Kaz0vS3jL5jzpV0nCd5zdC1PM2V1C/3J8yYFBeEmLHgXMWDobqUdv4u3qQ6LFlTbwebb/iGcTDiLh6nNxrxyHUke/w52um5ASKOHJywEVzwzExVqbJOlPdC/41fXe2Fd8B2Wkl5xVy4cdsC4/LcwSBCWOw4jIrphiu53PlWte6Ow/+G1rzHHfL+mwQ7SVESHNsZCcxSSSZMisFPjudTPM5LM3iftXuysx7EFTzDC3Z+sG4lvBGd0P18RCT+n3y35bBDXHF8Ui772S9KWu1+rxin91HkpiYb4rEpJMUPk7xXy45GbR+6TbGt3tzG+abnXcKrnpoiQUwWvOIHK3wNE1KbTxOTgc4ZNWI/mJ8bRqeIYDcm2eiPgxk6GNp5sKUQNwWDIH3KSh4NKkvbzl05ZDNa8HFBnk1Eka6QfH9S0hEyk+oS/5kINQbq8hJVpwnByhzY7B82I13UnddHwayBBR4wI8H5Xl83GLRq+wQXDBC/B4I52kglPmQqj373APLi15DZ2srIUY14UGL2wzcJFiOxLx6A3npJJC9inPE5Ar8/3Lv/zs/bqvep4hH7MwUnNyVEZH5wxJVGbkK+n3A5MyN0LXx0KZKeaPhWw2aDfx/hOE7fodCXcf86NxmDQab3/YhLQ447W2cXEVTsx3miwxXLOmnwKtl6WPfosbNZdJ0pvo2QqnB+JcxbXS9CcOaHa/PzZSHG0JBgpdG7Z81AHr81NwCGn6Vd837bGIeMOS2ac4fPMuY8VM4x8V1gILPPZlSN8aGtIyQ8WshQw/hrbIUDErlDKMftcCCyhfsoRit7JT4ZvnTbDIS1pxO+e1QZudX2Lrl3SbIP54fCki11feuD7timTax9JDyyDhSQhCTlEsZ3PFClQ/1AIPRx1E8hsJahzHodqlfnjSfjE0ZrZ7viWdYg7sMQ0yjTjlJomLV8yDR0v4SAvisl86G9y4AXx0MdHw6Mp72nKo5wpBS6QKr5jXtuRNB1Ilts9eePjXkcxORUlNQpucgtf5atKZw2Qckutfg3UG9CitHAgpdAOeb6TpUfS1UlAR95AMFLfcGmHuV+GcUFJ606lwRqT1Obgk6+hT1jWgcugPuZV0tqtssLT86yF4q+n2TloZ67NOfYz6IKaY/7XsWAP/AgEuKgpxT17h7RwjPcphqe4qu8M0yevKvNuEq+w7DckmJ7rL5k/MgcuvI2ireTQkUqUcKp7+Bg8nXkTI78JGIH0hW2r7Z19jZw1/hD+KlJiZDC0ftsfWsv5IjhKfquQyObq/aI2V+f2RliJWDmTH7vesCebyTmBixW7jocag+w3wq9c+ycXbo6QLupz4CjM99ktuDEq2zYNmy8tjBs87FmNTf3Z5VOidFws8zTzEdYrNcJUtpZg5Dt8GC5G4VucR86ZpNn11dubthcRmAsK+w5MO/PuHS20R3TfoT8xUf6EqpdszhV+fT5l++6DpAol0zcKJmW5bQtbvFysXnWLm1m9D3FXxt0Hd2S+ZDSyeAW3QK5Hs5XT1/vNS2K4YBG2cxMZBLkPS8A1w3NQFslTxppJTWCOm+1Z4HWwIDuK0lFprB4TV2wvvy19CdBNNo7ErBNTeBtm14lLsKyQ71xfyLb+WPhHHu4xGvF23DBSz8SpbH1DEHIBIW91VdrIxoYWV60oo1KZOdP9iOch2r3qoZwtj+sg2ZfP2Q5NGZzsspAbETsw5QkwfPsi08EgEjJ+LtIgo3uZp7ekC68mC4n3TZSa0FHtYt0DTwmz7mS88Zvbk6+82XKJzHtL1y3FwrPY58k0S7F5/1Fqhi3Er1NP7Pp0rwnegkGzgXMV1/MnJqAA4fHWtN2ytlbBSybDBdzO49Kde3WlZP8t1+XaahWHk0Od1ewM9apkPnTrTYcIJp2V9mcNfN5vaoEfqbMD0zEJRNC0ZhupCXF5Z9BinJj9Ihw05eMkwNlSwAR8c/Sdur31haF8X4tdQv7fjBQScpihCQv+8jUupxNC3zWBjpcDaL48j5EG0QblRvcJaicGvhPjeu6sdRRQ5rhmGz0FlrUD3J4Li3lfJH8mRien652Bjb4Xmt4T645V3Ii0+xWT+dp62qH5CGP+V6muQpqXzv174gEsZb5Rd3RzJqVo8brYIWl65GesdyheA7xTh/XfdfuaTgxjGR6fh2uXg2F9w3osbMkpHK9dPgIN9mxaQ1dUFzxj3DShstL7QpkMx1Rie0WpOdxMTOW06kkduNDzvsLYV77CmL7QHi+1tjL7mcbCOycaBD5HZTAjBScXzwpcmnw61lNr4OhJTNEhISoPX43LpsOEovDaCCgnX++ShbRO90jh4ukKXcQjyFWzIrtFCXOz0f1icUoEwN+F9J3SCLOFeeuj4/NBRNkIIT1dtQ2hTAnnZyRS+sHLbBLkyt0l/lvKDq/oXfn3Q24D1f0cf++fwxDE5AkKmmHOEmD5skBSPlxaZsOFTICM6yQTpXeLbQb/CytUO7lONXrnpe7rfYRUc8rojzxxh0Tcvd1r6wa6ABz5bIHZUIe/fxzOuw62il0lSBIpvHROfioD1L3Bn8W241/QyxMxO3/7DFY9xe+kTFG1fGFWmFufnk75cm/sUd9c8Rck+hVBppJj7eXrifTzY+hpVRxY38YjWt3Fq6G08OfIWX04ohhJdxfa8vf1v4PX5YDSaXU6UUIPa2N7xAkLuJmDIPelk99ubn0Ps8wT0vSd43OrpUZGxKTxHdnv900gMTkHP2w0ksd1b7xRS4tPQ4ZJ0+8fqHubfq39C2knofMN9UKitUXWfdP2N5pth5eWEkquaQqWQw9lBhdAogT5F5UnXVVDndkSeWcJJPH1RyGUIHjQbsvw+cBtlSrvjn7t5H/Fbt8OucCGgn8BnNilXLoI7sRuyz8sD7Ywey4Znrv8Oqyv+QNGvkFLXNOsZPWN/azsU9/2hKd4ccWVNx0ebP9s7S4AXR5BYuBviPzP1aKf3KUa2TcgFxBQcgDTvlnCyU/HZo/TFNWQkrBP+RIwHnYSNHtP6ereIQbBKeYB33uJwpfSMe1RPIO0NwtyNG4L083dObg+FLBLhVsclZe+S2hZyZQLS7I7BRi0wASytONvMSndzRR7+xHHQl4/3c3Ti9zkCSqaYc4SYPmyQ2uRkBMWli4h08z6ilm6GVsuTcKCQy+G00khhid1/GTEbT0AgT1HwETlybTJmS4n0u4CI7Zf5evqTkVupUGiXkaj/bPo5hB56qMuu+lfiARsrVDjWyzDom4POIfzCSyHJm0wGtZM1qp8RvIipXOt+GFF/hvG2aVqC7LwcUe1oc1AADkdbJdZX9kf002iefkPPOBVyQuMjRoW2q/rviHut86aVAd7lvNFwj5EeRdfjSfzCK7Sfv24e1FxtpEetK3wQqQlCPf1Xom1BVFtIcZCFsvnzPUhN0fDjJ6/bUn2LoNzEEiAOLQVQWei+k6dHCUchGb4cXwJF02XP2pB/LzRcGu9YRF65dZdWQ+502bN2FyGHOK1u/kIM8fTleJmtArWN6EucHA3uCB7W+nKu4gZodNZdBeSocb2HoY42SM9/OgWtjE7JHBQyGSqc/dpQT/So4JXHoaG2+QxKchTaY6SUaI5dRtjGI6DNHr0vVyhRdOcUfnPBc8n9DyLh6EkDvUuhVMFmwUzj4PzXAH9cFswDhI/SGhi32FDvcHAO0p7/IeQEJwVk64iEvisM9U5HRwHviD4l3DKr7LwQ3k5vngE8z/YEF210SJM5+ULbfCu/+aMxel9tC01SEH8K5/8inEohtPRyvn3qzvthXWiTY3RXqhxSiT7la6Rn+byrDk1qLG8nJwRTHVub0qeiKkCTRtQ54fNJce6FGIWRPuWcSA5laTy2MrkMyephiEcfw/xc/pY+9d/MPmXygf8Hf2CK2cKEnnDtD8T9eRtxZy9CPV/gb0b0/UFsU9Zq4bpGWEBfd5gqaXPOs208Xy+Z5EIuR+G9w3lP37tdN4tQVDioUfFQd8Sfe4eLQ06IInM55rbDFwdb4/XKx3iw6JopfQkcPMv7oOHepjj6zXk83GYMGq/vqHD7fPxJ+/e25xB6Pdg0ww84lBtZms9OtbvWUcS+iBVRZJpsrAvPGs7YUPQAkuOMuZz59jkOXS41hF1eO/gV3A2tLr2gvm9S0L1eCtzh1Xn3iEKC0rVcvzfCSW1tnp38gpy+UH3vV8L75pm3+O61QNunwknwWKlNImc6UhL17wgnzdMV10najGtdFzZGV2sYlZhh/FoZKl8QEsY/aT5P5Igsk2lReN8ovj6o+ySRuwFxkL03/sRfWccOGSnuXw44LJwLhIRAtozMJ6bzl9naQDt6IU+fkvnPENGb5C5eiOu+AHYP9kF5jYJ7mBbOoyiiG8+C858zYPWKTp0mdg2k5asPRY2fwJ3tB0XMLRH1Lq5wP8T79Ib389ZAnCl9iiYVXXCJQJ96UwNaranNm7amIXnJs74A3MPKA2Y2bcIk3EO42payKdP7UbZ3+fq/pU8lEH2qA5R2DSC3YHuzaPFgv6BDSAb8BAZOjkQg+WUAAidN45WFet6v/OIY/uN8CYcdDq6rZyBm4xlE+5vmkOUnznHIs308AkZuQMoTIfOKSaG8y/4jcaPpBqTGGK8BDYs/x6HK+a9xoso2cElmio8e0gL1bnfDsTJbJOMdU32ngD5Y4rZJkh5F2qTj83ZYb2aP1vdPp/teL9tl6OUtV8jQ/VkbrPSVTkKhslWg1+MWWJ+PgnmIp2/v64C2FxpglUR2J3o6d73caLj2C6zLu1O86eGA8qNLCRsHs5SYPPQAGh1vym8MjpTeLOHMxaHhHeEK+UyldRKi4VDrWm9h09Rlk4QzF/DF+a95elT03j9E75PiL+w/AjELtyDhD4nMYuDgvWEy4n6eCoRKOPKR3XnxXGDmt5Ali78NmiA36TfYLu0qyQCgby/hu61w3kB2awkvbIpk13M/3A/UkeANkx7WIqXTOVgf+kLSC5ts3OFfXYDn7YqSBADIVAgqeRGer8tI0qs4uRtCcp+Ae1gGSSrkJRDhuiVDL+wU276Ix7AM6VORtv684tcnubD2OgC5lWl+8xy5OLFBvzcCTDG/N1Sf7kGKshQSHgl3Vyc+SldGJfnZc4Rt3g76P6+Yx45HxJy1QIQ4LSO14bpqOt72nQ9NnDj1Hi2eebaPw5NW8yXTLlL7hf1H4UoNU2cYw9g4oPL5/jhRhk580vSn1k96YXfB9dL0KS3Q+H53bMpHuZyl3+/4vD025N0l+T5dHXZ/2SZDxUrj7xnQLkPFTPPo/6Y11uffKfZQ102y58u2WJ2bFLdYc6vUcvR40gpr8xIvWSwxtbsanf5oKqmY6WnHMi6oOrcaLjSlXNTi9/P1LAG1OhEv1z2U3DRV3N0eL6ceRfwD8aaKtuKVz3+NZz1XQBNpGlBDtydDEf8RCO3zMzTkJCixKfP2m4S4oaMy5KTbL5oL2eR+0oqVTCI/rYTN4o7SvGSDYqac0+LJ09cQ3WMfPA/U4SObiXcWQGjzE/A8V02Sl0xNJje8Ausr0opXK+MQUvI6PF+VkaQ3aeVyhOa+kaFiBtQ8PUqvWM3Hp5HbIEadcb1WmRvRVkvhkkAbEzIdVYTKeRTk1oJzZUYlTaP5v+vDp1vBWM8figBTzB+KWBY/f+bSnxj18zIkJArRqH4a2QsdmouTqnPR0Yi7/xChK1YbRqieu5D/d0Q/o71YX0mLs9vq6fyPrzuSvdnsuhUc8m4bj5h9NxG8+qTU2ozP94/Ei4knEXLqmXht5GT44nw/XG13CDFPI8T1MqDLy37YVm4jUsPFpypSZvVudcWe4ruRkiymRymsZGj3oB38CkjnU1Y7WKPDnWYZnpgd8juhzel6WJV7jyR9qsBXvqi3+Qusy79LSi+i5MCiqDC2BFZRQBCJJxqsqYY89b2xNq+phzkPBAf00l1l7/p8mziSFji0fdyRf/RoabGZgDYVDfQn5oprJU/kNXVX2VdqrJBoH6h8TrAzP2o+T5Ie9bn/SKQevYQwPzGvmi7ZfPwmA/sOIO64hAMUx8F+8Two/ZZC85yuks0KeeNO/A12fqPARUkELJFzSBi8Fc5bugKpUgFLZIjqsReehxqBD8BuXlRWCGl0CF7na0hmHpOrnKBtfAKy86Uk6VFadXGEFF4P7zdl+BjW5iXFqSEinWbBLZToU+LuE23G8PQp53jKx2xaT5uKKB09yjleCChiUv6KEhdpJ1x1p78KV7n+CoVtBchkXuIOKZpfYAgadx2DY1vnIJe3u+Qz7Jc5BwGmmLOxrBKTUvBV6+8wpE9rdG1TD6cv3sLQCYtwZMts5PbxMBm5Jj4er0ePgzZBOP0qnZ2hnDCZ/3fkwAnQpgqxm3m9wAHWPp6wnyo4+bztNAMaSlBsZLjApnR+uI8XrkuftpwPjieAGilALg1Kw+NbwQHrSo1VJvQp+t1nA6rAo5vAyzxRdqOJ8qPFqf6tbrxzl621Als/E5IE6Ast/A3vdDeEttz62XYz+hTQ6bnReUx0nc1x6BFgdKBan59CKhoLHcB7BBjjS1NO6PTgyCFH30DBE/2Pmfdwd7nxVCo4gBnfP9H/Bl4efmkcuy5+dt9XrWGrVmJ3o+N4dzPMSI8i5zqFDD2eCzbow/UPIj4gzqReZSVHi7vC+M/X3YeE0PSnWg5WNtaodUVXX3sz0uKS073PwdrRDlVPCDbq6zVXQ8tnD9HJDhxsCrqj9Dph/s9aLoAmvWzBwbVCPuSZ1hFJKRqE9posfBvp5ONQpQTsBgn4J3w7EmSCN3xb4OBQrw7QUqB/ySb3B8fJTOrxk/GWxXZJF7NTN4eEb7ca8HTc0NLkMptMFDE9iGsuFE9/SmBvyr8KaWH0dvY6Z3QCpOfJiY1rfp338k9I1vDX2SZezzIgqOQ1vm27pLWwDxE2t3whBzi5DEF5hM2GU/wEKBP3mlDPKCpdqJvAS6YEFfKkq+k+LXpfjkgbwQZN3tdc6mMT+pRCoUK4tZDRyi2lHjRpQUL7cleovX+HTOFmHI/uX50HTcHt+8LmmClmETw58hdMMWdjsdFpedAP83Hz6EpYWan4kTbp9j2vpLu2MdJoaLGJT0xD+ISJkFkroRopPiHTu9EjZwGODnD6yZiQIv30gwcugpWXM1wmS0dgCui2AsoCrvCdYlSK+vfJ+/f1wsuwLe+LzyeJT/Tv9gTg6fwrcK2THyUmGT2i3RytEJeowePVD/Bi5Z8o0LEIio0qL6JH3Zn7FM+3PsBnnYqhlAQ96tbE+3i87wVKDyps4hGtH9/Fsbfx6nAAvhhdGp9J0KNIwYZcDkaNGeVMPKb175/sdgGh96PQ8UZTyS+GnNBinieg061GBm6mp7M1wmNSkKbR4lCz00gIj0e7S9LvH2t8GFyKFg1ONJFs/1LLw9AkJaP6EYFPbF6utRWUVaVddP0rLndab4bMTY2Sq8T0IXr6Zc/lUHi78vQo2ji4O1gjOEq4pVEq5AgbPAectzNcfjR6FKfvJfmHCZDnzQXVN9+IO390HTi4GbLchcB1EH97dneOg7u6A1yeMkhsIFF/fzeUt3eBK1QfMRWN3v76juwfbIA6wB9J+ZojrpjRI53qiTtu92Q+VG8PIyZfN2jzd4edjdKQWYuecX43GTbRZxHjMQzxbsZkKvr2iZ9sHXsdMbl/NEk4oa93ixsEmfYBwhyl6VNOXE+oNC8QpjwrKZv3ok9pE5Dqcozf7JnTp0LCohAUEg5S0EwxS0Kc437JFHM2Ftl2/9NYt+13HNpozH7z7biFyJ/HByMHGvmbmtgYBK/eAHQX+Mia+UuQ8OSlzjYr8AMd61QD2gpBOCLGL0Pam7f8SUOg4ABuXZtCUb8KX//22yXQhOqzNwn23TxjOgPlC/L/ftl5MVITaNGm7FACgHStrS9/NlqNNL3DF232ORnKnR5gqL9Q3Q+aVCEDD511iEL05ZXewiJqo8SeoiuFU5ruJEZhGL+6alxwyWGMTvAC95ECgChQ/08jhehgse0gO6GuFgqZHE3uGTcTuwvvAIVp5HvnOCiVKrR6YFRYh0psJxSFuclkPP2rwXVj/ZFSFI/6LzspnQTBwcrZCrXOtecDqERceIdL/Q8b7OLUhK23A6ocErAP3v4Yzxec5d/VsbNg95kzSm8UTrAvl1xC8JbbulOYQDFyKZkbhZcJCj1oxmFEnHlocrvhWrMovH8Q+MpBY3Yj+n4APy+illHxbl0W+lSMEd8vQ/LbYJ3sBIC92jeEooUg+6gxc6CNCOf/rWeT5hnVE1zZkvxmKfWnH/l42IQ7OeDR/K3nLjLIVjl7EIiux1PbyD4qA9LGGzGwLzUAABM+SURBVM0rdsu7CsFtBPIcjyNdW+uL+87mPK2PRs4/IZchvN1BQ73X7/Wh5WUv1MtlCgQ3pqAeQqHoXvrLH34EcgXSGlzhN38paRp4PazIZ0UT3uegUFrjbeFLhvc935UVNlY6BDiZDUJ9rhjqXaJKGRNRkAgVjgizF5wn7XAGqoQhhrlTDwpVLoRbHdXVr4Eybb6Q9Ev3cSuUBRABYVPliFmQpWw0zI3/h6oMoiEEWLHHSDi4zIBM7mQYj/4fwaGRqNN+OFPMImRy5i+YYs7Gclu1+SAOn7qKnSuFK2kqZG+2t7XBpFHCySHxzm3EXzyH+BvXoZq+lP9d7KBRIrsjLTaOupCcwb0nSNhFOXiuncK/z9uczWxf5GLju1WIHPak5VwRasTQLbRvOIK330bgsouiegpKUebEABDnOXBfurR+uidlKjna3huA4233IeLPd6L3rRytUPVEF1DgjMQgsd3RPo8jqh5ojiMV9+k2DaZNuFXwRlW/mthXag/SksVe4gXa5OfpV4dLbIdWIuJzualV4dM6H46V3iIsqukKKfCGOvrSqQp+Ep6+HGrfEDYWl6oZlZShCY5D1UvkKEVmgd/EXtgw2oTvNaKQmeb0Kw4lDgtmiYfN5knarIseELi174j+ZN4DRW7bKKQ2jOgnQa3jBEdBKsmjKAWkaf9yOaD65Vfg+u+Q/04hQ83G91eqTO24VbC+uB7yG7+LE6ioVIgfsAHuZ8eACxLsq+mLXGGH0LY74HWmJ7j4N+J6G08E1d4Cn2uNoE2VSLDiUAqosRbKqxWg1Zr7M3BIcmqJKJ+f4BVUDhx/bW8q3HiXsXzqRre40rRfM60Ghwgn4eraNam0MRWq7in6NqLshHqntFJibwSqV+nqU8WxtOlbi1YJmDimlITKpjlU6maQq+uZjIMpZtFnkaN/wRRzNhbf+5yY0yLCEX/hHD8Leb1mQHw8Uk6dl3AI4mDVrCGS7z9F2mNxvGLSNXat6yDu2A1oo6QzBDm2r4nIPdfBJUllCAJcu1RD0IYb4sWNBkcOQ70r4dWqW5IOObSAFR1SEQ8WXZdOFM9xyDegHJ4vvS3ppUxdfPZNaTxZfE+kOKmOznaFh5TA/V+NIRHNRV/suxJ4sviuZPtyKzkK9i+OZ0uofwkvbCcl8nYthpfLhfCK5sU+rwPcm3yGN6sE+6F5cSjtCacvciNwjZi+RLLxbFcKKicrhPpdFo+PAzy6V0FqdAqiD/whqdjdu1ZF8u2nSH0qVmzUvn3bWkg5ew2aKLEHP81W3aIetKdOgEuRoD+R3bxhE+CcP+Tmik13MuVqtoLq2i5J2dNOJvWL9rC9Zwy9mR4f2gYllugG2yfrpJNAUJ7owj1h93qVZD0pN03BAVAFLpFOwCLjEOc2EHaxxtCgJv3LFEiw/xo2SUslHQE5uRqJVn2gTl0i2b9G4YkUeTuotUskvcy18vJIQRWoNVL1HJIUxFl3MdQrbNpDrvBhijkbr93/dmhMMf9bBD/i+3ob861jq6BSKfmeGnYejR7tG5jYmD/iEFjTDAGGQA5AgJ2Yc4CQPmCITDF/AFhZ/WhCYjIqNR6A7wd3Rpe/8crO6rGx/hgCDIHsgUBqmoZ3/mrUZQzvj0J0KZUy43gH2WPUbBT/DwGmmLP593Hywk2Qw5e+jB/WHZ1bEUWEFYYAQ4AhAFRqPNAQ54DwcHV2wLm9Roc8hlHOQ4Ap5hwgM41Gi6DQCHi6ORuutHPAsNkQGQIMAYYAQ+AfIMAU8z8Ajb3CEGAIMAQYAgyBj4UAU8wfC9ksaJcoUMT3/X/xs7NgGB/URVhENOxsbWCjtnqv9+h5ezsbqK3f7/n3ajQTHnrf+OXmXYVHxvC/cnNxzIRRZE4TKSmpiIyOg6e7s7RHfOZ0kyWtfMhcSIYRUTH8LZSTg12WjC8zOomNSwDFxXZxcsiM5lgb2RABppizoVDed0j+Ry9i/sodOLnDmD/2fd/N6udeBQZj4PfzEPCGglsAbZp8hYkjembopHLh2l0sXrsHb96GICk5FVUrFsf0sf15Jf2py/vGL9ePkxTA6i0HsWHHEURExcLWRo1rvws5gT9loY3dsg37sWTtHn4YZJtcPH0YyhQXAsn8v6JnDCydMRw1q5b5u8c/ev2HzuXS9Xv4bsIig222UtmiGPVNR5QsUuCjj/WfdkDx8r+fugLkd0KldPGCWDT1Oz65DSuWhQBTzDlQnqTk+o+agzfvQuHl4ZIjFPPXo+fwSnXa2P68B2mHAZMxcXgPNG9QTSQBOg2UqduXjxE+sHsLJCYlo13/n9CuWU306SQdsjKrxPgh8cv1Y5q7fDv2Hj6HgT1aonGdykhJTYW3h2tWDTnDfm7efYJuQ6bBb9GPKFX0M/y6ejcOnriE49vmQW6WQzp9I4+evebfI0WRXRTzh87l8o37CA2LwldVyyApKQU/z18P2kAtmzn8k8slowFQwKEd/qfht2gcf+P0zdj5KJDXB1PGSIdJzbYTYQP7WwSYYv5biLLfA6S46Ir35PmbWLX5QLZXzNGx8ajWfDA2Lh6HciUL84BOW+iHoJAILJo2VASwniY29fu+aN24Bl//44yVUCgUn3wRet/45fpJhYZHoVbbYUg/l+zyRdGG4cHTAKyaM5ofEsVcrt1uGB9prljhfJLDpPl0HDgZI77ugMnz1mPOxG+yxYn5n8wl/QTp9mns9N/w54nV2dY0RJvThrUqoX9XIUHIkdNXMWLSUtw9tTbHmyCyy99EdhkHU8zZRRL/YBy/n7yC2cu2ZnvF/OxlIFr0GofTuxbAw82Zn6nfzqPYd+SCSbjR9BDMW7Edq7ccQu9OjXklMXPRJvw2e1SGCuMfwPePXnmfaGzpGz5x7ga+m/ArOrWsg8fP38DaWoUWDaqhRYMv/1H/mfkShXd1cbLHuKHGpCUlavXK8BRMtwW9hs5Ajcql+dsMoulkF8X8oXMxx5GU8tMXgRl+j5mJ+z9ti/CmDR4pZyr3H79E+68n4aL/khxlI/+n8/8vvccUcw6Wdk5RzPprxvQLCCm45Rv2ZbipoKvG0T8v4+1olO7yy0olMXviN598AXqf+OXpP6lNu49j+q8beUVW5LM8ePT8NRav2YNfJgxE07pC4ohPVci8UKRgXpOEKLT4Uxx287HRNS8pPyqkjOmqOzsp5g+Zizne+tMy3RxUrVjiU4nj//ZLNvSStXubbJr0G97j2+bCx0ucDjJbToQN6r0QYIr5vWDKng/lFMWsX0DO7F5ocFT5fydm/dX3mvnfo3K5YiCb+pAfF6JQgdyYN0k6ZWVWSehDT8ykmLftO4n964UkEFTodEZ2zQU/D8mqYUv2Q4qWHL5+/K6boT6jE7P+mpvs/HY2av759TuOoFa1svzpX3+K+1QT+pC5pB8jORmSUv9pRE90aFH7Uw3/vfqljdC0sf3QoGZF/nl2Yn4v2HLkQ0wx50ixCYPOKYpZysY8Zf4GhIRFStqYz125g4Hfz8WFfYvh7GTPz5U8mhet2fPJvZk/NH654fnjqw0e6KREyKFtyXQhI9SnKmSXffTsFW8ioPL/bMzk6LVxlzG9Ij2/cNUuNKtfFc3qVeWvtz9l+ZC56Mept9FmR/u/FJZkY25U+wv06yKkAGU25k/5xX3cvpli/rj4fpTW6VorLU3Dp4QkutSRzbP5vLXZmc/cb9RsONrb8Tt+c69s4mX2Hj4LfTs34b2WA4PC0KDTKAzq2RJfd2uOxOQUDBo7Hw72tp/ca/bv4pebzyUmLgF1249Az/YN8U3Plrj76AW6DJrC23W7tP60oVWNnszjUKrYZ1i4aicOnbhs8Mq+dushZi3Zgrk/DUK+3F6ibzk7XWX/3VzWbT8MsveTBzoV8m8gh8KxQ7qgTvXyhrmRzZ3obNmxrNx0ADsPnOG9sm1trHn6IfPKzo6S+vdjYor532OY5S2Qk0rL3kJuZH0h2tHMH7/O8rG8b4cvXr3jFxKieFFp1ag6Jo3sxQd3iI6JR7UWg5E+DjidBvx2HgNRc6jQ9d23fdtkC5rR/4tfLjUXc84sKeTvh3T55Bsp2uARV3z5hv08xqSQfps90uA5f+riTd6EsHv1FBQpmCdbK+a/m8vspVtBZgg9f/zn+Rt4E4N5yc6n5/iEJN7Of/aykFqUONfEaqDAMKxYFgJMMVuWPLP9bCg9HfGZ7Wzf71RC9BxHBztYW6my1dw+NH45Udxo7tnxRJaUnIKIyBh4e7r9X/5ythJABoOxpLlkhDeZhlJT01hgkZzwQf7DMTLF/A+B+yevkb1Ro9Gg2Of54eNpDDDxKjAEdKLUR1Cia7k3b0Mlg2/8k37ZOwwBhsB/EwHaqNy+/xzPX71FckoqfL3deYdKMguxkn0RYIo5C2VDHq9UalQuheWzRhp61lNq7p1ex/9u0px12HHgNPQ/Z+EQWVcMAYaAhSBw484T/DD9N0OEQDplU0hYKtn5yt5C4P9X02CK+V/B92Evk2IumC8XngW8xfqFP6BimSJ8A+aKmRyMUtPSPjln98Nmx55mCDAEsgsCeg97Cs4ze8JA3kmMCplTfl29iw/0M6x/u+wyXDYOMwSYYs7CT4IU8/eDO2PP7+egVltj85LxfCg9c8VMp2VyGJo3aXAWjo51xRBgCFgKAhQudfv+U/h90yzk9RV71NPmnzy7J/yyBgXyeqNwgdygQCsh4VFYOOVbdij4xB8CU8xZKABSzBTMIa+vJ++hTDxWCtBgrphpR7v38PlsH2ozC6FjXTEEGAIfgECLnj/C18fjb+mFxI1+8CSAb5nWIoVCjimj+8LJMeekwfwAWHLMo0wxZ6Go9Iq5a5t6fHaemNh47FkzFVv3neTDNuptykwxZ6FQWFcMAQtDQJ+drUf7hvwN3f8rpJiJskiHBIoCx0r2QIAp5iyUQ3rFTI4Z3b+dhlnjBoDoD0wxZ6EgWFcMAQtHgIK/UJSwv0sJSYqZgstQSFJWsg8CTDFnoSzSK2bqdvCPC/hrpG5t64NCCrITcxYKg3XFELBgBOhGLj4hkb+R+7sTM1PM2e9DYIo5C2VirpgpqlWbvhP4KySiMTDFnIXCYF0xBCwYAQrfuWDlTsyfPMSQ9EI/XQobS3ETKHMbOzFnz4+AKeYslIu5YqauKdMQeUNSYYo5C4XBumIIWDACFEykdZ/xCHgTjMG9WuHLL0rxwY0ePHnFp1tt27QmT5diijl7fgRMMWehXEgxmycvoKhfjbuOMVHMi9bs5ilVJ3fMz8LRsa4YAgwBS0KAEqgsWr0Lm/ecMJlWnS/LYVCvViCOc8cBk1G8SH5mY85mgmeKOZsJhA2HIcAQYAhkJgKU4CM0PBrJKSnwcneBVTaLO5+Zc7WUtphithRJsnkwBBgCDAGGgEUgwBSzRYiRTYIhwBBgCDAELAUBppgtRZJsHgwBhgBDgCFgEQgwxWwRYmSTYAgwBBgCDAFLQYApZkuRJJsHQ4AhwBBgCFgEAkwxW4QY2SQYAgwBhgBDwFIQYIrZUiTJ5sEQYAgwBBgCFoEAU8wWIUY2CYYAQ4AhwBCwFASYYrYUSbJ5MAQYAgwBhoBFIMAUs0WIkU2CIcAQYAgwBCwFAaaYLUWSbB4MAYYAQ4AhYBEIMMVsEWJkk2AIMAQYAgwBS0GAKWZLkSSbB0OAIcAQYAhYBAJMMVuEGNkkGAIMAYYAQ8BSEGCK2VIkyebBEGAIMAQYAhaBAFPMFiFGNgmGAEOAIcAQsBQEmGK2FEmyeTAEGAIMAYaARSDAFLNFiJFNgiHAEGAIMAQsBQGmmC1FkmweDAGGAEOAIWARCDDFbBFiZJNgCDAEGAIMAUtBgClmS5EkmwdDgCHAEGAIWAQCTDFbhBjZJBgCDAGGAEPAUhBgitlSJMnmwRBgCDAEGAIWgQBTzBYhRjYJhgBDgCHAELAUBJhithRJsnkwBBgCDAGGgEUgwBSzRYiRTYIhwBBgCDAELAUBppgtRZJsHgwBhgBDgCFgEQgwxWwRYmSTYAgwBBgCDAFLQYApZkuRJJsHQ4AhwBBgCFgEAkwxW4QY2SQYAgwBhgBDwFIQYIrZUiTJ5sEQYAgwBBgCFoEAU8wWIUY2CYYAQ4AhwBCwFASYYrYUSbJ5MAQYAgwBhoBFIMAUs0WIkU2CIcAQYAgwBCwFAaaYLUWSbB4MAYYAQ4AhYBEIMMVsEWJkk2AIMAQYAgwBS0GAKWZLkSSbB0OAIcAQYAhYBAJMMVuEGNkkGAIMAYYAQ8BSEGCK2VIkyebBEGAIMAQYAhaBAFPMFiFGNgmGAEOAIcAQsBQEmGK2FEmyeTAEGAIMAYaARSDAFLNFiJFNgiHAEGAIMAQsBQGmmC1FkmweDAGGAEOAIWARCDDFbBFiZJNgCDAEGAIMAUtBgClmS5EkmwdDgCHAEGAIWAQCTDFbhBjZJBgCDAGGAEPAUhBgitlSJMnmwRBgCDAEGAIWgQBTzBYhRjYJhgBDgCHAELAUBJhithRJsnkwBBgCDAGGgEUgwBSzRYiRTYIhwBBgCDAELAUBppgtRZJsHgwBhgBDgCFgEQgwxWwRYmSTYAgwBBgCDAFLQeB/xqxmLTe5wAgAAAAASUVORK5CYII=",
       "text/html": [
        "