diff --git a/code/ARAX/ARAXQuery/Infer/scripts/infer_utilities.py b/code/ARAX/ARAXQuery/Infer/scripts/infer_utilities.py
index 5cac6db82..d2cf1053a 100644
--- a/code/ARAX/ARAXQuery/Infer/scripts/infer_utilities.py
+++ b/code/ARAX/ARAXQuery/Infer/scripts/infer_utilities.py
@@ -14,6 +14,7 @@ def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)
 from ARAX_expander import ARAXExpander
 from ARAX_resultify import ARAXResultify
 from ARAX_decorator import ARAXDecorator
+from biolink_helper import BiolinkHelper
 import traceback
 from collections import Counter
 from collections.abc import Hashable
@@ -43,7 +44,8 @@ def __init__(self):
         self.response = None
         self.message = None
         self.parameters = None
-        self.report_stats = True  
+        self.report_stats = True
+        self.bh = BiolinkHelper()
 
     def __get_formated_edge_key(self, edge: Edge, kp: str = 'infores:rtx-kg2') -> str:
         return f"{kp}:{edge.subject}-{edge.predicate}-{edge.object}"
@@ -129,8 +131,13 @@ def genrete_treat_subgraphs(self, response: ARAXResponse, top_drugs: pd.DataFram
                 self.response.envelope.message.knowledge_graph.nodes = {}
                 self.response.envelope.message.knowledge_graph.edges = {}
             # FIXME: is this the best way to be adding the node to the knowledge graph?
-            self.response.envelope.message.knowledge_graph.nodes[disease] = Node(name=disease_name, categories=[
-                'biolink:DiseaseOrPhenotypicFeature'])
+            #self.response.envelope.message.knowledge_graph.nodes[disease] = Node(name=disease_name, categories=[
+            #    'biolink:DiseaseOrPhenotypicFeature', 'biolink:Disease'])
+            categories_to_add = set()
+            categories_to_add.update(self.bh.get_ancestors('biolink:Disease'))
+            categories_to_add.update(list(synonymizer.get_normalizer_results(disease_curie)[disease_curie]['categories'].keys()))
+            categories_to_add = list(categories_to_add)
+            self.response.envelope.message.knowledge_graph.nodes[disease] = Node(name=disease_name, categories=categories_to_add)
             drug_qnode_key = response.envelope.message.query_graph.edges[qedge_id].subject
             disease_qnode_key = response.envelope.message.query_graph.edges[qedge_id].object
             self.response.envelope.message.knowledge_graph.nodes[disease].qnode_keys = [disease_qnode_key]
diff --git a/code/ARAX/test/test_ARAX_expand.py b/code/ARAX/test/test_ARAX_expand.py
index b9488a5db..eed50aba3 100644
--- a/code/ARAX/test/test_ARAX_expand.py
+++ b/code/ARAX/test/test_ARAX_expand.py
@@ -1083,5 +1083,67 @@ def test_xdtd_expand():
     nodes_by_qg_id, edges_by_qg_id = _run_query_and_do_standard_testing(json_query=query)
 
 
+def test_xdtd_different_categories():
+    query = {
+            "nodes": {
+                "disease": {
+                    "ids": ["MONDO:0004975"]
+                },
+                "chemical": {
+                    "categories": ["biolink:Drug"]
+                }
+            },
+            "edges": {
+                "t_edge": {
+                    "object": "disease",
+                    "subject": "chemical",
+                    "predicates": ["biolink:treats"],
+                    "knowledge_type": "inferred"
+                }
+            }
+        }
+    nodes_by_qg_id, edges_by_qg_id = _run_query_and_do_standard_testing(json_query=query)
+    query = {
+        "nodes": {
+            "disease": {
+                "ids": ["MONDO:0004975"],
+                "categories": ["biolink:Disease"]
+            },
+            "chemical": {
+                "categories": ["biolink:Drug"]
+            }
+        },
+        "edges": {
+            "t_edge": {
+                "object": "disease",
+                "subject": "chemical",
+                "predicates": ["biolink:treats"],
+                "knowledge_type": "inferred"
+            }
+        }
+    }
+    nodes_by_qg_id, edges_by_qg_id = _run_query_and_do_standard_testing(json_query=query)
+    query = {
+        "nodes": {
+            "disease": {
+                "ids": ["MONDO:0004975"],
+                "categories": ["biolink:DiseaseOrPhenotypicFeature"]
+            },
+            "chemical": {
+                "categories": ["biolink:ChemicalMixture"]
+            }
+        },
+        "edges": {
+            "t_edge": {
+                "object": "disease",
+                "subject": "chemical",
+                "predicates": ["biolink:treats"],
+                "knowledge_type": "inferred"
+            }
+        }
+    }
+    nodes_by_qg_id, edges_by_qg_id = _run_query_and_do_standard_testing(json_query=query)
+
+
 if __name__ == "__main__":
     pytest.main(['-v', 'test_ARAX_expand.py'])