Issues with setting up Groovy AST transformation in Gradle

Hi,
I need to run custom Groovy AST transformation that modifies bytecode using gradle groovy plugin. It worked for me when AST transformation classes are imported from external library as mentioned in gradle documentation:

 configurations { astTransformation }
 dependencies {
     astTransformation("org.mymodules:myast:1.0.0"))
 }
 tasks.withType(GroovyCompile).configureEach { 
     astTransformationClasspath.from(configurations.astTransformation)
 }

However I need to apply it also on a few classes in the original module where the AST transformation is implemented. I am not able to configure it correctly. I have noticed that original groovy AST transformations like @ToString are applied at that project but not the one I have implemented.

I have tried something like this (no other gradle.build modifications):

 tasks.withType(GroovyCompile).configureEach {
     astTransformationClasspath.from(sourceSets.main.compileClasspath)
 }

But from bytecode it is visible that the transformation was not applied.

Any idea how to set it up?

Many thanks,
Lukas