diff --git a/src/Tasks/MigrationTaskTrait.php b/src/Tasks/MigrationTaskTrait.php
index aa54bb9f..7e17cc86 100644
--- a/src/Tasks/MigrationTaskTrait.php
+++ b/src/Tasks/MigrationTaskTrait.php
@@ -19,8 +19,8 @@
 use SilverStripe\ORM\DataObjectSchema;
 use SilverStripe\ORM\DB;
 use SilverStripe\ORM\Queries\SQLUpdate;
-use SilverStripe\Versioned\ChangeSet;
-use SilverStripe\Versioned\ChangeSetItem;
+use SilverStripe\Versioned\Staged\ChangeSet;
+use SilverStripe\Versioned\Staged\ChangeSetItem;
 use SilverStripe\Versioned\Versioned;
 use Symfony\Component\Console\Command\Command;
 use Symfony\Component\Console\Input\InputInterface;