diff --git a/examples/benchmarks/TFT/tft.py b/examples/benchmarks/TFT/tft.py index 633a875c0f..3efffb6c54 100644 --- a/examples/benchmarks/TFT/tft.py +++ b/examples/benchmarks/TFT/tft.py @@ -317,6 +317,6 @@ def to_pickle(self, path: Union[Path, str]): for attr in drop_attrs: orig_attr[attr] = getattr(self, attr) setattr(self, attr, None) - super(TFTModel, self).to_pickle(path) + super(TFTModel, self).to_hdf(path, key="data") for attr in drop_attrs: setattr(self, attr, orig_attr[attr])