diff --git a/ServiceScan.SourceGenerator/DependencyInjectionGenerator.cs b/ServiceScan.SourceGenerator/DependencyInjectionGenerator.cs
index 9aa9dd7..a2dd531 100644
--- a/ServiceScan.SourceGenerator/DependencyInjectionGenerator.cs
+++ b/ServiceScan.SourceGenerator/DependencyInjectionGenerator.cs
@@ -135,7 +135,7 @@ private static DiagnosticModel<MethodWithAttributesModel> ParseMethodModel(Gener
         var attributeData = new AttributeModel[context.Attributes.Length];
         for (var i = 0; i < context.Attributes.Length; i++)
         {
-            attributeData[i] = AttributeModel.Create(context.Attributes[i], context.SemanticModel.Compilation);
+            attributeData[i] = AttributeModel.Create(context.Attributes[i], context.TargetSymbol.ContainingAssembly);
 
             if (!attributeData[i].HasSearchCriteria)
                 return Diagnostic.Create(MissingSearchCriteria, attributeData[i].Location);
diff --git a/ServiceScan.SourceGenerator/Model/AttributeModel.cs b/ServiceScan.SourceGenerator/Model/AttributeModel.cs
index 6a55758..ab26ebb 100644
--- a/ServiceScan.SourceGenerator/Model/AttributeModel.cs
+++ b/ServiceScan.SourceGenerator/Model/AttributeModel.cs
@@ -15,7 +15,7 @@ record AttributeModel(
 {
     public bool HasSearchCriteria => TypeNameFilter != null || AssignableToType != null;
 
-    public static AttributeModel Create(AttributeData attribute, Compilation compilation)
+    public static AttributeModel Create(AttributeData attribute, IAssemblySymbol currentAssembly)
     {
         var assemblyType = attribute.NamedArguments.FirstOrDefault(a => a.Key == "FromAssemblyOf").Value.Value as INamedTypeSymbol;
         var assignableTo = attribute.NamedArguments.FirstOrDefault(a => a.Key == "AssignableTo").Value.Value as INamedTypeSymbol;
@@ -34,7 +34,7 @@ public static AttributeModel Create(AttributeData attribute, Compilation compila
         };
 
         var assignableToTypeModel = assignableTo is not null ? TypeModel.Create(assignableTo) : null;
-        var registrations = GetRegistrationsFromAssembly(assemblyType, compilation, typeNameFilter, asSelf, asImplementedInterfaces, assignableToTypeModel, lifetime);
+        var registrations = GetRegistrationsFromAssembly(assemblyType, currentAssembly, typeNameFilter, asSelf, asImplementedInterfaces, assignableToTypeModel, lifetime);
 
         var syntax = attribute.ApplicationSyntaxReference.SyntaxTree;
         var textSpan = attribute.ApplicationSyntaxReference.Span;
@@ -47,7 +47,7 @@ public static AttributeModel Create(AttributeData attribute, Compilation compila
 
     private static EquatableArray<ServiceRegistrationModel>? GetRegistrationsFromAssembly(
         INamedTypeSymbol? fromAssemblyOf,
-        Compilation compilation,
+        IAssemblySymbol currentAssembly,
         string? typeNameFilter,
         bool asSelf,
         bool asImplementedInterfaces,
@@ -58,7 +58,7 @@ public static AttributeModel Create(AttributeData attribute, Compilation compila
             return null;
 
         // if user specifies FromAssemblyOf = typeof(SomeType), but SomeType is from the same assembly as the method with the attribute
-        if (SymbolEqualityComparer.Default.Equals(fromAssemblyOf.ContainingAssembly, compilation.Assembly))
+        if (SymbolEqualityComparer.Default.Equals(fromAssemblyOf.ContainingAssembly, currentAssembly))
             return null;
 
         var types = fromAssemblyOf.ContainingAssembly