M1 chip causing problems with PyTorch

[katdapre]

Hi,

I'm on a Mac with an M1 chip and haven't been able to get PyTorch to work with it (including the instructions for this on the seisbench readme).

I can train models using the seisbench notebook 03a to train phasenet successfully by just commenting out the line
model.to_preferred_device(verbose=True)
and this makes everything run slowly but smoothly.

If I don't comment it out, it shows
Model device: mps:0
and later on when I run the training cell that iterates through epochs, it will cause the error:
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.
which I think implies that PyTorch has identified the GPU but the float64 format just isn't compatible with it (this is a known M1 issue, e.g. this apple dev issue)

However this means that when running train.py here, I get the same error, and I can't stop it from trying to use GPU because it's identified that MPS is available but then the float64 data isn't compatible.

I fully understand and respect that Apple's hardware is not your problem! I am just wondering if there's any way of turning off the move to GPU (like commenting out that line in the notebook) while running the full training workflow as found in this repository? Perhaps something I can pass to trainer.fit?

Alternatively, I know the notebook03a has that line in augmentations
sbg.ChangeDtype(np.float32), # as expected by the pytorch model
and I'm wondering if there's somewhere else I should be doing this (either in the notebook, or in the train.py workflow) to avoid the error.

If I do manage to get this to work I will definitely be documenting it to help out other M1 users, because it's causing me a world of pain!

Thanks very much,
Kat

[yetinam]

Hi @katdapre ,

thanks a lot for this comprehensive issue report. That's indeed an intruiging issue. I have no personal experience with Apple Silicon, but given we have quite a few SeisBench users with Macbooks, it's worth figuring out the underlying problem.

My suspicion is that the error occurs when trying to move the labels to the MPS device, because the label vector is not converted to float32. To try if this is indeed the problem, could you add a sbg.ChangeDtype(np.float32, key="y") augmentation at the end of the augmentations list? If this doesn't fix the problem, I'd ask you to pasted the full error stack trace, so I can have a more detailed look. Please also let me know if it works, so I can adjust the notebook accordingly.

For running the training code in this repository, the GPU is configured through the trainer_args in the json configs. If you remove it there, training will be executed on CPU. You can also adapt the parameters for training with MPS (https://lightning.ai/docs/pytorch/stable/accelerators/mps_basic.html#run-on-apple-silicon-gpus). However, I'm not totally certain if the pytorch lighning version pinned in this repository already suggests MPS, so you might need to updated it. In case the fix for the notebook works, you'll have to add the same ChangeDtype augmentations to the get_augmentations functions in models.py.

[katdapre]

Hi @yetinam

Thanks for the rapid response! Adding that line in notebook 03a has in fact allowed the notebook to run on GPU, which is very exciting.

In this repository, I have attached the couple of files I've changed:

models.py, which I added the sbg.ChangeDtype(np.float32, key="y") line to everywhere I could see an existing np.float32 line. It's possible that I missed some places where it needs to be added, though, either in this file or somewhere else, because I'm still getting the same error (see below).

The config file I'm using: test.json, which is sending to MPS I think in line with how the PyTorch lightning instructions that you linked recommended. If I delete the accelerator/gpus lines entirely, I still get exactly the same error as below - it seems like the move_data_to_device line 612 is being triggered anyway?

I should mention that requirements.txt didn't work for me when building the environment (couldn't seem to solve dependencies even when I just listed the packages and no version requirements), so I just used the working environment that I had for the tutorial notebooks and added the missing packages to it, which seemed to work fine. My environment details are here:
pip_environment.txt but given that this error is one I had with the notebooks earlier, it doesn't seem like a versioning issue to me.

Full error code below:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File ~/Documents/ML_python/pick-benchmark/benchmark/train.py:215
    213 if args.test_run:
    214     experiment_name = experiment_name + "_test"
--> 215 train(config, experiment_name, test_run=args.test_run)
    217 running_time = str(
    218     datetime.timedelta(seconds=time.perf_counter() - code_start_time)
    219 )
    220 print(f"Running time: {running_time}")

File ~/Documents/ML_python/pick-benchmark/benchmark/train.py:86, in train(config, experiment_name, test_run)
     75 ## Uncomment the following 2 lines to enable
     76 # device_stats = DeviceStatsMonitor()
     77 # callbacks.append(device_stats)
     79 trainer = pl.Trainer(
     80     default_root_dir=default_root_dir,
     81     logger=loggers,
     82     callbacks=callbacks,
     83     **config.get("trainer_args", {}),
     84 )
---> 86 trainer.fit(model, train_loader, dev_loader)

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/trainer.py:584, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path, weights_only)
    582 self.training = True
    583 self.should_stop = False
--> 584 call._call_and_handle_interrupt(
    585     self,
    586     self._fit_impl,
    587     model,
    588     train_dataloaders,
    589     val_dataloaders,
    590     datamodule,
    591     ckpt_path,
    592     weights_only,
    593 )

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/call.py:49, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
     47     if trainer.strategy.launcher is not None:
     48         return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
---> 49     return trainer_fn(*args, **kwargs)
     51 except _TunerExitException:
     52     _call_teardown_hook(trainer)

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/trainer.py:630, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path, weights_only)
    623     download_model_from_registry(ckpt_path, self)
    624 ckpt_path = self._checkpoint_connector._select_ckpt_path(
    625     self.state.fn,
    626     ckpt_path,
    627     model_provided=True,
    628     model_connected=self.lightning_module is not None,
    629 )
--> 630 self._run(model, ckpt_path=ckpt_path, weights_only=weights_only)
    632 assert self.state.stopped
    633 self.training = False

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/trainer.py:1079, in Trainer._run(self, model, ckpt_path, weights_only)
   1074 self._signal_connector.register_signal_handlers()
   1076 # ----------------------------
   1077 # RUN THE TRAINER
   1078 # ----------------------------
-> 1079 results = self._run_stage()
   1081 # ----------------------------
   1082 # POST-Training CLEAN UP
   1083 # ----------------------------
   1084 log.debug(f"{self.__class__.__name__}: trainer tearing down")

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/trainer.py:1121, in Trainer._run_stage(self)
   1119 if self.training:
   1120     with isolate_rng():
-> 1121         self._run_sanity_check()
   1122     with torch.autograd.set_detect_anomaly(self._detect_anomaly):
   1123         self.fit_loop.run()

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/trainer.py:1150, in Trainer._run_sanity_check(self)
   1147 call._call_callback_hooks(self, "on_sanity_check_start")
   1149 # run eval step
-> 1150 val_loop.run()
   1152 call._call_callback_hooks(self, "on_sanity_check_end")
   1154 # reset logger connector

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/loops/utilities.py:179, in _no_grad_context.<locals>._decorator(self, *args, **kwargs)
    177     context_manager = torch.no_grad
    178 with context_manager():
--> 179     return loop_run(self, *args, **kwargs)

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/loops/evaluation_loop.py:146, in _EvaluationLoop.run(self)
    144     self.batch_progress.is_last_batch = data_fetcher.done
    145     # run step hooks
--> 146     self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
    147 except StopIteration:
    148     # this needs to wrap the `*_step` call too (not just `next`) for `dataloader_iter` support
    149     break

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/loops/evaluation_loop.py:415, in _EvaluationLoop._evaluation_step(self, batch, batch_idx, dataloader_idx, dataloader_iter)
    413     batch = trainer.precision_plugin.convert_input(batch)
    414     batch = trainer.lightning_module._on_before_batch_transfer(batch, dataloader_idx=dataloader_idx)
--> 415     batch = call._call_strategy_hook(trainer, "batch_to_device", batch, dataloader_idx=dataloader_idx)
    417 # the `_step` methods don't take a batch_idx when `dataloader_iter` is used, but all other hooks still do,
    418 # so we need different kwargs
    419 hook_kwargs = self._build_kwargs(
    420     batch, batch_idx, dataloader_idx if self._is_sequential and self.num_dataloaders > 1 else None
    421 )

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/call.py:329, in _call_strategy_hook(trainer, hook_name, *args, **kwargs)
    326     return None
    328 with trainer.profiler.profile(f"[Strategy]{trainer.strategy.__class__.__name__}.{hook_name}"):
--> 329     output = fn(*args, **kwargs)
    331 # restore current_fx when nested context
    332 pl_module._current_fx_name = prev_fx_name

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/strategies/strategy.py:278, in Strategy.batch_to_device(self, batch, device, dataloader_idx)
    276 device = device or self.root_device
    277 if model is not None:
--> 278     return model._apply_batch_transfer_handler(batch, device=device, dataloader_idx=dataloader_idx)
    279 return move_data_to_device(batch, device)

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/core/module.py:362, in LightningModule._apply_batch_transfer_handler(self, batch, device, dataloader_idx)
    358 def _apply_batch_transfer_handler(
    359     self, batch: Any, device: Optional[torch.device] = None, dataloader_idx: int = 0
    360 ) -> Any:
    361     device = device or self.device
--> 362     batch = self._call_batch_hook("transfer_batch_to_device", batch, device, dataloader_idx)
    363     batch = self._call_batch_hook("on_after_batch_transfer", batch, dataloader_idx)
    364     return batch

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/core/module.py:351, in LightningModule._call_batch_hook(self, hook_name, *args)
    348     else:
    349         trainer_method = call._call_lightning_datamodule_hook
--> 351     return trainer_method(trainer, hook_name, *args)
    352 hook = getattr(self, hook_name)
    353 return hook(*args)

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/call.py:177, in _call_lightning_module_hook(trainer, hook_name, pl_module, *args, **kwargs)
    174 pl_module._current_fx_name = hook_name
    176 with trainer.profiler.profile(f"[LightningModule]{pl_module.__class__.__name__}.{hook_name}"):
--> 177     output = fn(*args, **kwargs)
    179 # restore current_fx when nested context
    180 pl_module._current_fx_name = prev_fx_name

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/core/hooks.py:612, in DataHooks.transfer_batch_to_device(self, batch, device, dataloader_idx)
    565 def transfer_batch_to_device(self, batch: Any, device: torch.device, dataloader_idx: int) -> Any:
    566     """Override this hook if your :class:`~torch.utils.data.DataLoader` returns tensors wrapped in a custom data
    567     structure.
    568 
   (...)    610 
    611     """
--> 612     return move_data_to_device(batch, device)

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/lightning_fabric/utilities/apply_func.py:110, in move_data_to_device(batch, device)
    107     # user wrongly implemented the `_TransferableDataType` and forgot to return `self`.
    108     return data
--> 110 return apply_to_collection(batch, dtype=_TransferableDataType, function=batch_to)

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/lightning_utilities/core/apply_func.py:72, in apply_to_collection(data, dtype, function, wrong_dtype, include_none, allow_frozen, *args, **kwargs)
     70     return tuple(function(x, *args, **kwargs) for x in data)
     71 if data.__class__ is dict and all(isinstance(x, dtype) for x in data.values()):  # 1d homogeneous dict
---> 72     return {k: function(v, *args, **kwargs) for k, v in data.items()}
     73 # slow path for everything else
     74 return _apply_to_collection_slow(
     75     data,
     76     dtype,
   (...)     82     **kwargs,
     83 )

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/lightning_fabric/utilities/apply_func.py:104, in move_data_to_device.<locals>.batch_to(data)
    102 if isinstance(data, Tensor) and isinstance(device, torch.device) and device.type not in _BLOCKING_DEVICE_TYPES:
    103     kwargs["non_blocking"] = True
--> 104 data_output = data.to(device, **kwargs)
    105 if data_output is not None:
    106     return data_output

TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

[yetinam]

Hi @katdapre ,
sorry for the late response. I agree, it's most likely not a versioning error, but a similar error to before. From what I've seen in the code, ChangeDtype seems to be added at all necessary spots. Potentially even a few too many, so you might see crashes from that. Why it still fails is unclear to me. What I'd recommend to do is some manual debugging. For example, you could create the data generators, add the augmentations and then manually query an example (e.g., train_generator[0]). Then you can check the dtype of each entry and see if they are all correct.

[yetinam]

Btw, I've incorporated your finding for the example notebook in the main SeisBench repository (seisbench/seisbench@4d563c6). Thanks again!

[katdapre]

Hi @yetinam

I'm getting back to this following AGU and Christmas and seem to have working training code on the M1 MPS! I can do a full 100-epoch Phasenet run on my own dataset (~7k traces) in 2.5 hours, which seems pretty good. Once I'm sure everything works, I'm thinking of making some kind of instruction manual for getting seisbench running on M1?

I'm now looking at making eval.py work. I'm running into a key error (see below) which is pretty reminiscent of previous CPU vs GPU issues - it seems to me like the model produced with train.py was on a different architecture to where I'm trying to load it in the eval step. I don't know if maybe there's somewhere I should be sending the model to GPU manually or similar?

I've attached:
The config file for the model run
train.py, eval.py, models.py, and generate_eval_targets.py, which are the only scripts I've modified in benchmark

RIS_Phasenet.json
eval.py
generate_eval_targets.py
models.py
train.py

Let me know what you think, or if there's anything else I can send over.

Cheers,
Kat

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File ~/Documents/ML_python/pick-benchmark/benchmark/eval.py:230
    225 parser.add_argument(
    226     "--sampling_rate", type=float, help="Overwrites the sampling rate in the data"
    227 )
    228 args = parser.parse_args()
--> 230 main(
    231     args.weights,
    232     args.targets,
    233     args.sets,
    234     batchsize=args.batchsize,
    235     num_workers=args.num_workers,
    236     sampling_rate=args.sampling_rate,
    237 )
    238 running_time = str(
    239     datetime.timedelta(seconds=time.perf_counter() - code_start_time)
    240 )
    241 print(f"Running time: {running_time}")

File ~/Documents/ML_python/pick-benchmark/benchmark/eval.py:149, in main(weights, targets, sets, batchsize, num_workers, sampling_rate)
    144 loader = DataLoader(
    145     generator, batch_size=batchsize, shuffle=False, num_workers=num_workers
    146 )
    147 trainer = pl.Trainer(accelerator="mps", devices=1)
--> 149 predictions = trainer.predict(model, loader)
    151 # Merge batches
    152 merged_predictions = []

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/trainer.py:941, in Trainer.predict(self, model, dataloaders, datamodule, return_predictions, ckpt_path, weights_only)
    939 self.state.status = TrainerStatus.RUNNING
    940 self.predicting = True
--> 941 return call._call_and_handle_interrupt(
    942     self,
    943     self._predict_impl,
    944     model,
    945     dataloaders,
    946     datamodule,
    947     return_predictions,
    948     ckpt_path,
    949     weights_only,
    950 )

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/call.py:49, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
     47     if trainer.strategy.launcher is not None:
     48         return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
---> 49     return trainer_fn(*args, **kwargs)
     51 except _TunerExitException:
     52     _call_teardown_hook(trainer)

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/trainer.py:990, in Trainer._predict_impl(self, model, dataloaders, datamodule, return_predictions, ckpt_path, weights_only)
    986     download_model_from_registry(ckpt_path, self)
    987 ckpt_path = self._checkpoint_connector._select_ckpt_path(
    988     self.state.fn, ckpt_path, model_provided=model_provided, model_connected=self.lightning_module is not None
    989 )
--> 990 results = self._run(model, ckpt_path=ckpt_path, weights_only=weights_only)
    992 assert self.state.stopped
    993 self.predicting = False

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/trainer.py:1079, in Trainer._run(self, model, ckpt_path, weights_only)
   1074 self._signal_connector.register_signal_handlers()
   1076 # ----------------------------
   1077 # RUN THE TRAINER
   1078 # ----------------------------
-> 1079 results = self._run_stage()
   1081 # ----------------------------
   1082 # POST-Training CLEAN UP
   1083 # ----------------------------
   1084 log.debug(f"{self.__class__.__name__}: trainer tearing down")

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/trainer/trainer.py:1118, in Trainer._run_stage(self)
   1116     return self._evaluation_loop.run()
   1117 if self.predicting:
-> 1118     return self.predict_loop.run()
   1119 if self.training:
   1120     with isolate_rng():

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/loops/utilities.py:179, in _no_grad_context.<locals>._decorator(self, *args, **kwargs)
    177     context_manager = torch.no_grad
    178 with context_manager():
--> 179     return loop_run(self, *args, **kwargs)

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/loops/prediction_loop.py:122, in _PredictionLoop.run(self)
    120 else:
    121     dataloader_iter = None
--> 122     batch, batch_idx, dataloader_idx = next(data_fetcher)
    123 self.batch_progress.is_last_batch = data_fetcher.done
    124 # run step hooks

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/loops/fetchers.py:134, in _PrefetchDataFetcher.__next__(self)
    131         self.done = not self.batches
    132 elif not self.done:
    133     # this will run only when no pre-fetching was done.
--> 134     batch = super().__next__()
    135 else:
    136     # the iterator is empty
    137     raise StopIteration

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/loops/fetchers.py:61, in _DataFetcher.__next__(self)
     59 self._start_profiler()
     60 try:
---> 61     batch = next(self.iterator)
     62 except StopIteration:
     63     self.done = True

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/utilities/combined_loader.py:341, in CombinedLoader.__next__(self)
    339 def __next__(self) -> _ITERATOR_RETURN:
    340     assert self._iterator is not None
--> 341     out = next(self._iterator)
    342     if isinstance(self._iterator, _Sequential):
    343         return out

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/pytorch_lightning/utilities/combined_loader.py:142, in _Sequential.__next__(self)
    139             raise StopIteration
    141 try:
--> 142     out = next(self.iterators[0])
    143 except StopIteration:
    144     # try the next iterator
    145     self._use_next_iterator()

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/torch/utils/data/dataloader.py:732, in _BaseDataLoaderIter.__next__(self)
    729 if self._sampler_iter is None:
    730     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    731     self._reset()  # type: ignore[call-arg]
--> 732 data = self._next_data()
    733 self._num_yielded += 1
    734 if (
    735     self._dataset_kind == _DatasetKind.Iterable
    736     and self._IterableDataset_len_called is not None
    737     and self._num_yielded > self._IterableDataset_len_called
    738 ):

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/torch/utils/data/dataloader.py:1506, in _MultiProcessingDataLoaderIter._next_data(self)
   1504 worker_id = self._task_info.pop(idx)[0]
   1505 self._rcvd_idx += 1
-> 1506 return self._process_data(data, worker_id)

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/torch/utils/data/dataloader.py:1541, in _MultiProcessingDataLoaderIter._process_data(self, data, worker_idx)
   1539 self._try_put_index()
   1540 if isinstance(data, ExceptionWrapper):
-> 1541     data.reraise()
   1542 return data

File ~/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/torch/_utils.py:769, in ExceptionWrapper.reraise(self)
    765 except Exception:
    766     # If the exception takes multiple arguments or otherwise can't
    767     # be constructed, don't try to instantiate since we don't know how to
    768     raise RuntimeError(msg) from None
--> 769 raise exception

KeyError: Caught KeyError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/Users/katdapre/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/torch/utils/data/_utils/worker.py", line 349, in _worker_loop
    data = fetcher.fetch(index)  # type: ignore[possibly-undefined]
  File "/Users/katdapre/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
            ~~~~~~~~~~~~^^^^^
  File "/Users/katdapre/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/seisbench/generate/generator.py", line 78, in __getitem__
    func(state_dict)
    ~~~~^^^^^^^^^^^^
  File "/Users/katdapre/.local/share/virtualenvs/ML_python-AKGzaeOV/lib/python3.13/site-packages/seisbench/generate/augmentation.py", line 277, in __call__
    x, metadata = state_dict[self.key[0]]
                  ~~~~~~~~~~^^^^^^^^^^^^^
KeyError: 'y'