diff --git a/DJT.EntityFrameworkCore/DJT.EntityFrameworkCore.csproj b/DJT.EntityFrameworkCore/DJT.EntityFrameworkCore.csproj
index a401983..a1208b7 100644
--- a/DJT.EntityFrameworkCore/DJT.EntityFrameworkCore.csproj
+++ b/DJT.EntityFrameworkCore/DJT.EntityFrameworkCore.csproj
@@ -16,7 +16,7 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.EntityFrameworkCore" Version="6.0.0" />
+    <PackageReference Include="Microsoft.EntityFrameworkCore" Version="6.0.28" />
   </ItemGroup>
 
 </Project>
diff --git a/DJT.EntityFrameworkCore/DbContextExtensions.cs b/DJT.EntityFrameworkCore/DbContextExtensions.cs
index 95ac59a..3d422b8 100644
--- a/DJT.EntityFrameworkCore/DbContextExtensions.cs
+++ b/DJT.EntityFrameworkCore/DbContextExtensions.cs
@@ -21,24 +21,23 @@ public static class DbContextExtensions
         /// If entities are internal, the assembly attribute <see cref="InternalsVisibleToAttribute"/> may
         /// need specifying.
         /// </summary>
-        /// <param name="dbContext"></param>
         /// <param name="modelBuilder"></param>
-        public static void UseSelfDefiningModels(this DbContext dbContext, ModelBuilder modelBuilder)
+        public static void UseSelfDefiningModels(this ModelBuilder modelBuilder)
         {
-            Type dbContextType = dbContext.GetType();
-            foreach (var prop in dbContextType.GetProperties())
+            var entities = modelBuilder.Model.GetEntityTypes();
+            foreach (var entity in entities)
             {
-                var propType = prop.PropertyType;
-                if (propType.IsGenericType && propType.GetGenericTypeDefinition() == typeof(DbSet<>))
+                modelBuilder.Entity(entity.Name, builder =>
                 {
-                    var itemType = propType.GenericTypeArguments[0];//.GetGenericTypeDefinition().GetGenericArguments()[0];
-                    var method = itemType.GetMethod("OnModelCreating");
-                    if (typeof(ISelfDefine).IsAssignableFrom(itemType) && method != null)
+                    Type itemType = entity.ClrType;
+
+                    var method = itemType.GetMethod(nameof(ISelfDefine.OnModelCreating));
+                    if (typeof(ISelfDefine).IsAssignableFrom(itemType) && method is not null)
                     {
-                        //If the "OnModelCreating(ModelBuilder modelBuilder)" method exists, run it with the given ModelBuilder
-                        method.Invoke(Activator.CreateInstance(itemType, false), new object[] { modelBuilder });
+                        method.Invoke(Activator.CreateInstance(itemType, false),
+                            new object[] { modelBuilder });
                     }
-                }
+                });
             }
         }
     }