diff --git a/coverage.xml b/coverage.xml
index c2e75b3..f7f3de8 100644
--- a/coverage.xml
+++ b/coverage.xml
@@ -1,12 +1,12 @@
 <?xml version="1.0" ?>
-<coverage version="7.4.3" timestamp="1717636032945" lines-valid="315" lines-covered="303" line-rate="0.9619" branches-valid="96" branches-covered="82" branch-rate="0.8542" complexity="0">
+<coverage version="7.4.3" timestamp="1717638283749" lines-valid="319" lines-covered="306" line-rate="0.9592" branches-valid="96" branches-covered="82" branch-rate="0.8542" complexity="0">
 	<!-- Generated by coverage.py: https://coverage.readthedocs.io/en/7.4.3 -->
 	<!-- Based on https://raw.githubusercontent.com/cobertura/web/master/htdocs/xml/coverage-04.dtd -->
 	<sources>
 		<source>/home/jvivian/covid19-drDFM/covid19_drdfm</source>
 	</sources>
 	<packages>
-		<package name="." line-rate="0.9619" branch-rate="0.8542" complexity="0">
+		<package name="." line-rate="0.9592" branch-rate="0.8542" complexity="0">
 			<classes>
 				<class name="cli.py" filename="cli.py" complexity="0" line-rate="0.8857" branch-rate="0.5">
 					<methods/>
@@ -114,7 +114,7 @@
 						<line number="139" hits="1"/>
 					</lines>
 				</class>
-				<class name="dfm.py" filename="dfm.py" complexity="0" line-rate="0.9701" branch-rate="0.9">
+				<class name="dfm.py" filename="dfm.py" complexity="0" line-rate="0.9577" branch-rate="0.9">
 					<methods/>
 					<lines>
 						<line number="3" hits="1"/>
@@ -151,8 +151,9 @@
 						<line number="75" hits="0"/>
 						<line number="77" hits="1"/>
 						<line number="92" hits="1"/>
-						<line number="94" hits="1"/>
-						<line number="95" hits="1" branch="true" condition-coverage="100% (2/2)"/>
+						<line number="93" hits="1"/>
+						<line number="94" hits="1" branch="true" condition-coverage="100% (2/2)"/>
+						<line number="95" hits="1"/>
 						<line number="96" hits="1"/>
 						<line number="97" hits="1"/>
 						<line number="98" hits="1"/>
@@ -164,26 +165,29 @@
 						<line number="104" hits="1"/>
 						<line number="105" hits="1"/>
 						<line number="106" hits="1"/>
-						<line number="107" hits="1"/>
 						<line number="109" hits="1"/>
 						<line number="110" hits="1"/>
+						<line number="111" hits="0"/>
 						<line number="112" hits="1"/>
 						<line number="113" hits="1"/>
-						<line number="115" hits="1"/>
-						<line number="122" hits="1" branch="true" condition-coverage="100% (2/2)"/>
-						<line number="123" hits="1" branch="true" condition-coverage="100% (2/2)"/>
-						<line number="124" hits="1"/>
-						<line number="126" hits="1"/>
-						<line number="133" hits="1" branch="true" condition-coverage="50% (1/2)" missing-branches="134"/>
-						<line number="134" hits="0"/>
-						<line number="135" hits="1" branch="true" condition-coverage="100% (2/2)"/>
-						<line number="138" hits="1"/>
-						<line number="150" hits="1"/>
-						<line number="151" hits="1"/>
-						<line number="152" hits="1" branch="true" condition-coverage="100% (2/2)"/>
-						<line number="153" hits="1" branch="true" condition-coverage="50% (1/2)" missing-branches="155"/>
+						<line number="114" hits="1"/>
+						<line number="116" hits="1"/>
+						<line number="117" hits="1"/>
+						<line number="119" hits="1"/>
+						<line number="126" hits="1" branch="true" condition-coverage="100% (2/2)"/>
+						<line number="127" hits="1" branch="true" condition-coverage="100% (2/2)"/>
+						<line number="128" hits="1"/>
+						<line number="130" hits="1"/>
+						<line number="137" hits="1" branch="true" condition-coverage="50% (1/2)" missing-branches="138"/>
+						<line number="138" hits="0"/>
+						<line number="139" hits="1" branch="true" condition-coverage="100% (2/2)"/>
+						<line number="142" hits="1"/>
 						<line number="154" hits="1"/>
 						<line number="155" hits="1"/>
+						<line number="156" hits="1" branch="true" condition-coverage="100% (2/2)"/>
+						<line number="157" hits="1" branch="true" condition-coverage="50% (1/2)" missing-branches="159"/>
+						<line number="158" hits="1"/>
+						<line number="159" hits="1"/>
 					</lines>
 				</class>
 				<class name="io.py" filename="io.py" complexity="0" line-rate="1" branch-rate="1">
diff --git a/covid19_drdfm/dfm.py b/covid19_drdfm/dfm.py
index 4934999..c3a17e5 100644
--- a/covid19_drdfm/dfm.py
+++ b/covid19_drdfm/dfm.py
@@ -90,7 +90,6 @@ def run(self, maxiter=10_000, global_multiplier=1, columns: Optional[list[str]]
         - Exception: If an error occurs during model fitting.
         """
         self.outdir.mkdir(exist_ok=True)
-        # self.ad.obs = self.ad.obs.drop(columns="Time")
         print(f"{len(self.batches)} batches to run")
         for batch_name, batch in track(list(self.batches.items())):
             data = DataProcessor(batch, global_multiplier, maxiter).process(columns)
@@ -106,7 +105,12 @@ def run(self, maxiter=10_000, global_multiplier=1, columns: Optional[list[str]]
             result = Result(batch_name, res, model, filtered_factors)
             result.write(self.outdir)
             # self.ad.uns["factors"] = result.factors.drop(columns="Time")
-            self.ad.write(self.outdir / batch_name / "data.h5ad")
+            # TODO: Fix this. Tests need this present but the dashboard doesn't
+            try:
+                self.ad.obs = self.ad.obs.drop(columns="Time")
+                self.ad.write(self.outdir / batch_name / "data.h5ad")
+            except:
+                pass
             self.results.append(result)
         # TODO: Concat factors across batch variables
         print("All runs completed!")