diff --git a/rules-tests/CodeQuality/Rector/Class_/YieldDataProviderRector/Fixture/remove_return_non_array_doc.php.inc b/rules-tests/CodeQuality/Rector/Class_/YieldDataProviderRector/Fixture/remove_return_non_array_doc.php.inc
new file mode 100644
index 00000000..04eede44
--- /dev/null
+++ b/rules-tests/CodeQuality/Rector/Class_/YieldDataProviderRector/Fixture/remove_return_non_array_doc.php.inc
@@ -0,0 +1,48 @@
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\YieldDataProviderRector\Fixture;
+
+use PHPUnit\Framework\TestCase;
+
+final class RemoveReturnNonArrayDoc extends TestCase
+{
+    #[\PHPUnit\Framework\Attributes\DataProvider('dataProvider')]
+    public function test(string $filePath): void
+    {
+    }
+
+    /**
+     * @return mixed
+     */
+    public static function dataProvider()
+    {
+        return [
+            ['<?php implode(\'\', $foo, );', '<?php implode($foo, );'],
+            ['<?php implode(\'\', $foo, );', '<?php implode($foo, );']
+        ];
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\YieldDataProviderRector\Fixture;
+
+use PHPUnit\Framework\TestCase;
+
+final class RemoveReturnNonArrayDoc extends TestCase
+{
+    #[\PHPUnit\Framework\Attributes\DataProvider('dataProvider')]
+    public function test(string $filePath): void
+    {
+    }
+
+    public static function dataProvider(): \Iterator
+    {
+        yield ['<?php implode(\'\', $foo, );', '<?php implode($foo, );'];
+        yield ['<?php implode(\'\', $foo, );', '<?php implode($foo, );'];
+    }
+}
+
+?>
diff --git a/rules/CodeQuality/Rector/Class_/YieldDataProviderRector.php b/rules/CodeQuality/Rector/Class_/YieldDataProviderRector.php
index 5263408b..51710228 100644
--- a/rules/CodeQuality/Rector/Class_/YieldDataProviderRector.php
+++ b/rules/CodeQuality/Rector/Class_/YieldDataProviderRector.php
@@ -15,6 +15,7 @@
 use PHPStan\Type\Generic\GenericObjectType;
 use Rector\BetterPhpDocParser\PhpDocInfo\PhpDocInfoFactory;
 use Rector\BetterPhpDocParser\PhpDocManipulator\PhpDocTypeChanger;
+use Rector\Comments\NodeDocBlock\DocBlockUpdater;
 use Rector\DeadCode\NodeAnalyzer\IsClassMethodUsedAnalyzer;
 use Rector\NodeTypeResolver\Node\AttributeKey;
 use Rector\PhpParser\NodeTransformer;
@@ -40,6 +41,7 @@ public function __construct(
         private readonly PhpDocInfoFactory $phpDocInfoFactory,
         private readonly IsClassMethodUsedAnalyzer $isClassMethodUsedAnalyzer,
         private readonly PhpDocTypeChanger $phpDocTypeChanger,
+        private readonly DocBlockUpdater $docBlockUpdater
     ) {
     }
 
@@ -184,6 +186,7 @@ private function removeReturnTag(ClassMethod $classMethod): void
             );
         } else {
             $phpDocInfo->removeByType(ReturnTagValueNode::class);
+            $this->docBlockUpdater->updateRefactoredNodeWithPhpDocInfo($classMethod);
         }
     }
 }