diff --git a/src/ifcopenshell-python/ifcopenshell/api/root/reassign_class.py b/src/ifcopenshell-python/ifcopenshell/api/root/reassign_class.py
index 3c93eeee80..aef53b625f 100644
--- a/src/ifcopenshell-python/ifcopenshell/api/root/reassign_class.py
+++ b/src/ifcopenshell-python/ifcopenshell/api/root/reassign_class.py
@@ -17,6 +17,7 @@
 # along with IfcOpenShell.  If not, see <http://www.gnu.org/licenses/>.
 
 import ifcopenshell
+import ifcopenshell.util.type
 import ifcopenshell.util.schema
 import ifcopenshell.util.element
 
@@ -79,12 +80,12 @@ class Usecase:
                 element.PredefinedType = "USERDEFINED"
                 element.ObjectType = self.settings["predefined_type"]
         if element.is_a("IfcTypeProduct"):
-            for typed_element in ifcopenshell.util.element.get_types(element) or []:
-                ifc_class = self.settings["ifc_class"].removesuffix("Type").removesuffix("Style")
+            for occurrence in ifcopenshell.util.element.get_types(element) or []:
+                ifc_class = ifcopenshell.util.type.get_applicable_entities(self.settings["ifc_class"])[0]
                 ifcopenshell.api.run(
                     "root.reassign_class",
                     self.file,
-                    product=typed_element,
+                    product=occurrence,
                     ifc_class=ifc_class,
                     predefined_type= self.settings["predefined_type"]
                 )
diff --git a/src/ifcopenshell-python/test/api/pset/test_edit_pset.py b/src/ifcopenshell-python/test/api/pset/test_edit_pset.py
index 9d308fd559..d7c8fdf025 100644
--- a/src/ifcopenshell-python/test/api/pset/test_edit_pset.py
+++ b/src/ifcopenshell-python/test/api/pset/test_edit_pset.py
@@ -95,7 +95,6 @@ class TestEditPset(test.bootstrap.IFC4):
         element = ifcopenshell.api.run("root.create_entity", self.file, ifc_class="IfcWall")
         pset = ifcopenshell.api.run("pset.add_pset", self.file, product=element, name="Pset_WallCommon")
         ifcopenshell.api.run("pset.edit_pset", self.file, pset=pset, properties={"Reference": "Foo"})
-        print("test")
         ifcopenshell.api.run("pset.edit_pset", self.file, pset=pset, properties={"Reference": None}, should_purge=True)
         pset = element.IsDefinedBy[0].RelatingPropertyDefinition
         assert len(pset.HasProperties) == 0