import Union from typing

Author: yaoyi92

src/atomate2/forcefields/jobs.py

@@ -5,7 +5,7 @@
 import logging
 import warnings
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
 from ase.io import Trajectory as AseTrajectory
 from ase.units import GPa as _GPa_to_eV_per_A3

src/atomate2/forcefields/schemas.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from emmet.core.utils import ValueEnum
 from emmet.core.vasp.calculation import StoreTrajectoryOption
@@ -249,7 +249,7 @@ def from_ase_task_doc(
             meta_class = ForceFieldStructureTaskDocument
             k = "structure"
             if relax_cell := getattr(ase_task_doc, "relax_cell", None):
-                kwargs.update({"relax_cell": relax_cell})
+                task_document_kwargs.update({"relax_cell": relax_cell})
             task_document_kwargs.update(structure=ase_task_doc.mol_or_struct)
         elif isinstance(ase_task_doc.mol_or_struct, Molecule):
             meta_class = ForceFieldMoleculeTaskDocument