diff --git a/android/src/main/kotlin/io/customer/customer_io/CustomerIoPlugin.kt b/android/src/main/kotlin/io/customer/customer_io/CustomerIoPlugin.kt
index 63a59aa..7286a47 100644
--- a/android/src/main/kotlin/io/customer/customer_io/CustomerIoPlugin.kt
+++ b/android/src/main/kotlin/io/customer/customer_io/CustomerIoPlugin.kt
@@ -184,11 +184,10 @@ class CustomerIoPlugin : FlutterPlugin, MethodCallHandler, ActivityAware {
     }
 
     private fun registerDeviceToken(params: Map<String, Any>) {
-        // TODO: Fix registerDeviceToken implementation
-        /*
-        val token = params.getString(Keys.Tracking.TOKEN)
+        val token = requireNotNull(params.getAsTypeOrNull<String>(Keys.Tracking.TOKEN)) {
+            "Device token is missing in params: $params"
+        }
         CustomerIO.instance().registerDeviceToken(token)
-         */
     }
 
     private fun trackMetric(params: Map<String, Any>) {
diff --git a/ios/Classes/SwiftCustomerIoPlugin.swift b/ios/Classes/SwiftCustomerIoPlugin.swift
index 78ae4c0..3f44517 100644
--- a/ios/Classes/SwiftCustomerIoPlugin.swift
+++ b/ios/Classes/SwiftCustomerIoPlugin.swift
@@ -133,20 +133,22 @@ public class SwiftCustomerIoPlugin: NSObject, FlutterPlugin {
     
     private func setProfileAttributes(params : Dictionary<String, AnyHashable>){
         guard let attributes = params[Keys.Tracking.traits] as? Dictionary<String, AnyHashable>
-        else { return }
+        else {
+            logger.error("Missing attributes in: \(params) for key: \(Keys.Tracking.traits)")
+            return
+        }
+        
         CustomerIO.shared.profileAttributes = attributes
     }
     
     private func registerDeviceToken(params : Dictionary<String, AnyHashable>){
-        // TODO: Fix registerDeviceToken implementation
-        /*
-         guard let token = params[Keys.Tracking.token] as? String
-         else {
-         return
-         }
-         
-         CustomerIO.shared.registerDeviceToken(token)
-         */
+        guard let token = params[Keys.Tracking.token] as? String
+        else {
+            logger.error("Missing token in: \(params) for key: \(Keys.Tracking.token)")
+            return
+        }
+        
+        CustomerIO.shared.registerDeviceToken(token)
     }
     
     private func trackMetric(params : Dictionary<String, AnyHashable>){