diff --git a/src/ActionsGridFieldItemRequest.php b/src/ActionsGridFieldItemRequest.php
index ce9fafc..fa94530 100644
--- a/src/ActionsGridFieldItemRequest.php
+++ b/src/ActionsGridFieldItemRequest.php
@@ -28,7 +28,6 @@
 use ReflectionObject;
 use SilverStripe\Admin\ModelAdmin;
 use SilverStripe\Core\Extension;
-use SilverStripe\ORM\DataObjectInterface;
 use SilverStripe\View\ViewableData;
 
 /**
@@ -536,10 +535,10 @@ public function getPublicEditLinkForAdjacentRecord(int $offset): ?string
 
     /**
      * @param FieldList $actions
-     * @param DataObjectInterface $record
+     * @param ViewableData $record
      * @return void
      */
-    public function addSaveAndClose(FieldList $actions, DataObjectInterface $record)
+    public function addSaveAndClose(FieldList $actions, ViewableData $record)
     {
         if (!$this->checkCan($record)) {
             return;
@@ -567,10 +566,10 @@ public function addSaveAndClose(FieldList $actions, DataObjectInterface $record)
     /**
      * New and existing records have different classes
      *
-     * @param DataObjectInterface $record
+     * @param ViewableData $record
      * @return string
      */
-    protected function getBtnClassForRecord(DataObjectInterface $record)
+    protected function getBtnClassForRecord(ViewableData $record)
     {
         if ($record->ID) {
             return 'btn-outline-primary';