diff --git a/src/Control/ElementFormController.php b/src/Control/ElementFormController.php
index 1abd2c9..a19a4b2 100644
--- a/src/Control/ElementFormController.php
+++ b/src/Control/ElementFormController.php
@@ -2,13 +2,13 @@
 
 namespace DNADesign\ElementalUserForms\Control;
 
+use DNADesign\Elemental\Models\BaseElement;
 use DNADesign\Elemental\Controllers\ElementController;
 use SilverStripe\Control\Controller;
 use SilverStripe\Control\Director;
 use SilverStripe\Core\Injector\Injector;
 use SilverStripe\UserForms\Control\UserDefinedFormController;
 use SilverStripe\UserForms\Form\UserForm;
-use SilverStripe\CMS\Model\SiteTree;
 
 class ElementFormController extends ElementController
 {
@@ -61,7 +61,7 @@ public function finished()
             return null;
         }
 
-        while (!$page instanceof SiteTree) {
+        while ($page instanceof BaseElement) {
             $page = $page->getPage();
 
             if ($page === null) {