diff --git a/Model/SimpleProtect/SimpleProtect.php b/Model/SimpleProtect/SimpleProtect.php
index c470941..f272025 100644
--- a/Model/SimpleProtect/SimpleProtect.php
+++ b/Model/SimpleProtect/SimpleProtect.php
@@ -73,7 +73,7 @@ class SimpleProtect extends OrigSimpleProtect
     /**
      * Checkout session object
      *
-     * @var \Magento\Checkout\Model\Session\Proxy
+     * @var \Magento\Checkout\Model\Session
      */
     protected $checkoutSession;
 
@@ -89,13 +89,13 @@ class SimpleProtect extends OrigSimpleProtect
      *
      * @param \Payone\Core\Model\SimpleProtect\ProtectFunnel     $protectFunnel
      * @param \Magento\Framework\App\ResourceConnection          $resource
-     * @param \Magento\Checkout\Model\Session\Proxy              $checkoutSession
+     * @param \Magento\Checkout\Model\Session                    $checkoutSession
      * @param \Magento\Framework\App\Config\ScopeConfigInterface $scopeConfig
      */
     public function __construct(
         \Payone\Core\Model\SimpleProtect\ProtectFunnel $protectFunnel,
         \Magento\Framework\App\ResourceConnection $resource,
-        \Magento\Checkout\Model\Session\Proxy $checkoutSession,
+        \Magento\Checkout\Model\Session $checkoutSession,
         \Magento\Framework\App\Config\ScopeConfigInterface $scopeConfig
     ) {
         parent::__construct($protectFunnel);
diff --git a/etc/di.xml b/etc/di.xml
index ed47ff4..e9f2bdf 100644
--- a/etc/di.xml
+++ b/etc/di.xml
@@ -26,4 +26,9 @@
 -->
 <config xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:noNamespaceSchemaLocation="urn:magento:framework:ObjectManager/etc/config.xsd">
     <preference for="Payone\Core\Model\SimpleProtect\SimpleProtect" type="Payone\SimpleProtect\Model\SimpleProtect\SimpleProtect" /><!-- Use your own Simple Protect implementation instead of empty default implementation -->
+    <type name="Payone\SimpleProtect\Model\SimpleProtect\SimpleProtect">
+        <arguments>
+            <argument name="checkoutSession" xsi:type="object">Magento\Checkout\Model\Session\Proxy</argument>
+        </arguments>
+    </type>
 </config>