Update element_array_ephys/export/nwb/nwb.py

Co-authored-by: Kabilar Gunalan <kabilar@datajoint.com>

Author: bendichter

element_array_ephys/export/nwb/nwb.py

@@ -197,6 +197,7 @@ def create_units_table(
 
     units_table = pynwb.misc.Units(name=name, description=desc)
     # add additional columns to the units table
+    units_query = ephys.CuratedClustering.Unit() & session_key
     for additional_attribute in ["cluster_quality_label", "spike_depths"]:
         units_table.add_column(
             name=units_query.heading.attributes[additional_attribute].name,