diff --git a/py_gen/templates/_ofclass.py b/py_gen/templates/_ofclass.py
index 8c21bb3..4c5358e 100644
--- a/py_gen/templates/_ofclass.py
+++ b/py_gen/templates/_ofclass.py
@@ -1,5 +1,5 @@
-:: from py_gen.codegen import Member, LengthMember, TypeMember
-:: normal_members = [m for m in ofclass.members if type(m) == Member]
+:: from loxi_ir import *
+:: normal_members = [m for m in ofclass.members if type(m) == OFDataMember]
 class ${ofclass.pyname}(${superclass}):
 :: for m in ofclass.type_members:
     ${m.name} = ${m.value}