diff --git a/atomicfu-gradle-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/gradle/AtomicFUGradlePlugin.kt b/atomicfu-gradle-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/gradle/AtomicFUGradlePlugin.kt
index eb4d858e..dcc74bef 100644
--- a/atomicfu-gradle-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/gradle/AtomicFUGradlePlugin.kt
+++ b/atomicfu-gradle-plugin/src/main/kotlin/kotlinx/atomicfu/plugin/gradle/AtomicFUGradlePlugin.kt
@@ -337,8 +337,9 @@ private fun Project.configureTransformationForTarget(target: KotlinTarget) {
                 it.destinationDirectory.value(originalDestinationDirectory)
             }
         }
-        val originalClassesDirs: FileCollection =
-            project.files(classesDirs.from.toTypedArray()).filter { it.exists() }
+        val originalClassesDirs: FileCollection = project.objects.fileCollection()
+            .from(compilationTask.flatMap { it.destinationDirectory })
+            .from({ project.files(classesDirs.from).filter { it.exists() } })
         originalDirsByCompilation[compilation] = originalClassesDirs
         val transformedClassesDir = project.layout.buildDirectory
             .dir("classes/atomicfu/${target.name}/${compilation.name}")
diff --git a/integration-testing/examples/jvm-sample/build.gradle.kts b/integration-testing/examples/jvm-sample/build.gradle.kts
index 2b3f2d5d..26cf262e 100644
--- a/integration-testing/examples/jvm-sample/build.gradle.kts
+++ b/integration-testing/examples/jvm-sample/build.gradle.kts
@@ -14,6 +14,7 @@ group = "kotlinx.atomicfu.examples"
 version = "DUMMY_VERSION"
 
 plugins {
+    application
     kotlin("jvm") version libs.versions.kotlinVersion.get()
     `maven-publish`
 }
@@ -37,6 +38,10 @@ tasks.compileKotlin {
     }
 }
 
+application {
+    mainClass.set("org.example.MainKt")
+}
+
 publishing {
     repositories {
         /**
diff --git a/integration-testing/examples/jvm-sample/src/main/java/JavaClass.java b/integration-testing/examples/jvm-sample/src/main/java/JavaClass.java
index 5c6c713d..a5ea52f9 100644
--- a/integration-testing/examples/jvm-sample/src/main/java/JavaClass.java
+++ b/integration-testing/examples/jvm-sample/src/main/java/JavaClass.java
@@ -1,5 +1,3 @@
-package org.example;
-
 import java.util.concurrent.atomic.AtomicLongFieldUpdater;
 
 public class JavaClass {