include Union in forcefield/md.py

Author: yaoyi92

src/atomate2/forcefields/md.py

@@ -4,7 +4,7 @@
 
 import warnings
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
 from jobflow import job
 from monty.dev import deprecated