diff --git a/MVMCoreUI/Atomic/Atoms/TextFields/DigitEntryField.swift b/MVMCoreUI/Atomic/Atoms/TextFields/DigitEntryField.swift
index 977447ff..01305a95 100644
--- a/MVMCoreUI/Atomic/Atoms/TextFields/DigitEntryField.swift
+++ b/MVMCoreUI/Atomic/Atoms/TextFields/DigitEntryField.swift
@@ -334,17 +334,13 @@ import UIKit
             setAsSecureTextEntry(entryType == .secure || entryType == .password)
         }
         
-        for digitBox in digitBoxes {
-            
-            if let observingTextFieldelegate = delegateObject?.observingTextFieldDelegate {
-                digitBox.digitField.inputAccessoryView = UIToolbar.getToolbarWithDoneButton(delegate: observingTextFieldelegate,
-                                                                                            action: #selector(observingTextFieldelegate.dismissFieldInput))
-            } else {
-                digitBox.digitField.inputAccessoryView = UIToolbar.getToolbarWithDoneButton(delegate: self,
-                                                                                            action: #selector(dismissFieldInput))
-            }
-        }
+        let observingDelegate = delegateObject?.observingTextFieldDelegate ?? self
         
+        digitBoxes.forEach {
+            $0.digitField.inputAccessoryView = UIToolbar.getToolbarWithDoneButton(delegate: observingDelegate,
+                                                                                  action: #selector(observingDelegate.dismissFieldInput))
+        }
+            
         super.set(with: model, delegateObject, additionalData)
     }