diff --git a/liger_iris_pipeline/assign_wcs/assign_wcs.py b/liger_iris_pipeline/assign_wcs/assign_wcs.py
index 35910e2..7f1a617 100644
--- a/liger_iris_pipeline/assign_wcs/assign_wcs.py
+++ b/liger_iris_pipeline/assign_wcs/assign_wcs.py
@@ -5,6 +5,7 @@
 from astropy import coordinates as coord
 from astropy import units as u
 from gwcs import coordinate_frames as cf
+from gwcs import WCS
 
 log = logging.getLogger(__name__)
 log.setLevel(logging.DEBUG)
diff --git a/liger_iris_pipeline/assign_wcs/assign_wcs_step.py b/liger_iris_pipeline/assign_wcs/assign_wcs_step.py
index cdd85f2..0acc556 100755
--- a/liger_iris_pipeline/assign_wcs/assign_wcs_step.py
+++ b/liger_iris_pipeline/assign_wcs/assign_wcs_step.py
@@ -24,31 +24,23 @@ class AssignWCSStep(LigerIRISStep):
 
     def process(self, input, *args, **kwargs):
         reference_file_names = {}
-        if isinstance(input, str):
-            input_model = datamodels.open(input)
-        else:
-            input_model = input
-
-        # If input type is not supported, log warning, set to 'skipped', exit
-        if not (isinstance(input_model, ImagerModel)):
-            log.warning("Input dataset type is not supported.")
-            log.warning("assign_wcs expects ImageModel as input.")
-            log.warning("Skipping assign_wcs step.")
-            result = input_model.copy()
-            self.status = "SKIPPED"
-        else:
-            # Get reference files
-            for reftype in self.reference_file_types:
-                reffile = self.get_reference_file(input_model, reftype)
-                reference_file_names[reftype] = reffile if reffile else ""
-            log.debug(f"reference files used in assign_wcs: {reference_file_names}")
-
-            # Assign wcs
-            result = load_wcs(input_model, reference_file_names)
-            self.status = "COMPLETE"
-
-        # Close model if opened manually
-        if isinstance(input, str):
-            input_model.close()
+        with datamodels.open(input) as input_model:
+            # If input type is not supported, log warning, set to 'skipped', exit
+            if not (isinstance(input_model, ImagerModel)):
+                log.warning("Input dataset type is not supported.")
+                log.warning("assign_wcs expects ImageModel as input.")
+                log.warning("Skipping assign_wcs step.")
+                result = input_model.copy()
+                self.status = "SKIPPED"
+            else:
+                # Get reference files
+                for reftype in self.reference_file_types:
+                    reffile = self.get_reference_file(input_model, reftype)
+                    reference_file_names[reftype] = reffile if reffile else ""
+                log.debug(f"reference files used in assign_wcs: {reference_file_names}")
+
+                # Assign wcs
+                result = load_wcs(input_model, reference_file_names)
+                self.status = "COMPLETE"
 
         return result