diff --git a/source/MaterialXCore/Definition.cpp b/source/MaterialXCore/Definition.cpp
index 6b45baf195..cf89f93de6 100644
--- a/source/MaterialXCore/Definition.cpp
+++ b/source/MaterialXCore/Definition.cpp
@@ -185,16 +185,26 @@ StringVec TargetDef::getMatchingTargets() const
 
 vector<UnitDefPtr> UnitTypeDef::getUnitDefs() const
 {
-    const auto datalibrary = getDocument()->hasDataLibrary() ? getDocument()->getRegisteredDataLibrary() : getDocument();
-
     vector<UnitDefPtr> unitDefs;
-    for (UnitDefPtr unitDef : datalibrary->getChildrenOfType<UnitDef>())
+    for (UnitDefPtr unitDef : getDocument()->getChildrenOfType<UnitDef>())
     {
         if (unitDef->getUnitType() == _name)
         {
             unitDefs.push_back(unitDef);
         }
     }
+    // Gather unitdefs from Data library
+    if (getDocument()->hasDataLibrary())
+    {
+        for (UnitDefPtr unitDef : getDocument()->getRegisteredDataLibrary()->getChildrenOfType<UnitDef>())
+        {
+            if (unitDef->getUnitType() == _name)
+            {
+                unitDefs.push_back(unitDef);
+            }
+        }
+        
+    }
     return unitDefs;
 }
 
diff --git a/source/MaterialXCore/Document.h b/source/MaterialXCore/Document.h
index 6fc9523cca..6c0fbb6149 100644
--- a/source/MaterialXCore/Document.h
+++ b/source/MaterialXCore/Document.h
@@ -554,7 +554,7 @@ class MX_CORE_API Document : public GraphElement
     /// Return a vector of all Member elements in the TypeDef.
     vector<UnitDefPtr> getUnitDefs() const
     {
-        return getChildrenOfType<UnitDef>();
+        return hasDataLibrary() ? getChildrenOfType<UnitDef>(getRegisteredDataLibrary()) : getChildrenOfType<UnitDef>();
     }
 
     /// Remove the UnitDef, if any, with the given name.
@@ -585,7 +585,7 @@ class MX_CORE_API Document : public GraphElement
     /// Return a vector of all UnitTypeDef elements in the document.
     vector<UnitTypeDefPtr> getUnitTypeDefs() const
     {
-        return getChildrenOfType<UnitTypeDef>();
+        return hasDataLibrary() ? getChildrenOfType<UnitTypeDef>(getRegisteredDataLibrary()) : getChildrenOfType<UnitTypeDef>();
     }
 
     /// Remove the UnitTypeDef, if any, with the given name.