diff --git a/docs/derived_fields.md b/docs/derived_fields.md
index b834c268..cd954b84 100644
--- a/docs/derived_fields.md
+++ b/docs/derived_fields.md
@@ -86,18 +86,20 @@ def GroupDistance3D(arrs, snap=None):
         groupid = groupid.magnitude
         boxsize *= snap.ureg("code_length")
     pos_cat = snap.data["Group"]["GroupPos"][groupid]
-    dist3 = pos_part-pos_cat
+    dist3 = (pos_part-pos_cat)
+    dist3, u = dist3.magnitude, dist3.units
     dist3 = da.where(dist3>boxsize/2.0, boxsize-dist3, dist3)
     dist3 = da.where(dist3<=-boxsize/2.0, boxsize+dist3, dist3) # PBC
-    return dist3
+    return dist3 * u
 
 @fielddefs.register_field("all")
 def GroupDistance(arrs, snap=None):
     import dask.array as da
     dist3 = arrs["GroupDistance3D"]
+    dist3, u = dist3.magnitude, dist3.units
     dist = da.sqrt((dist3**2).sum(axis=1))
     dist = da.where(arrs["GroupID"]==-1, np.nan, dist) # set unbound gas to nan
-    return dist
+    return dist * u
 ```
 
 1. We define a list of field containers that we want to add particles to.
@@ -117,5 +119,4 @@ In above example, we now have the following fields available:
 gas = ds.data["PartType0"]
 print(gas["Volume"])
 print(gas["GroupDistance"])
-print(gas["Subhalo"])
 ```
diff --git a/tests/test_docs.py b/tests/test_docs.py
index 29bf7b32..67a50f67 100644
--- a/tests/test_docs.py
+++ b/tests/test_docs.py
@@ -47,8 +47,8 @@ def evaluate(self):
         # evaluate all at once (for now)
         code = "\n".join(self.codeblocks)
         print("Evaluating code:")
-        print(code)
-        exec(code)
+        print('"""' + code + '"""')
+        exec(code, globals())  # need to use global() to allow using imports
 
 
 def get_docfiles():