use the remembered device when switching

modules/two_factor_authentication/app/controllers/two_factor_authentication/authentication_controller.rb

@@ -84,24 +84,26 @@ def successful_2fa_transmission(service, transmit)
     ##
     # Create a token service for the current user
     # with an optional override to use a non-default channel
-    def otp_service(user, use_channel: nil, use_device: nil)
+    def otp_service(user, use_channel: nil, use_device: remembered_device)
       session[:two_factor_authentication_device_id] = use_device.try(:id)
       ::TwoFactorAuthentication::TokenService.new user:, use_channel:, use_device:
     end
 
     ##
     # Get the used device for verification
     def otp_service_for_verification(user)
-      use_device =
-        if session[:two_factor_authentication_device_id]
-          user.otp_devices.find(session[:two_factor_authentication_device_id])
-        end
-      otp_service(user, use_device:)
+      otp_service(user, use_device: remembered_device)
     rescue ActiveRecord::RecordNotFound
       render_404
       false
     end
 
+    def remembered_device
+      if session[:two_factor_authentication_device_id]
+        user.otp_devices.find(session[:two_factor_authentication_device_id])
+      end
+    end
+
     ##
     # Detect overridden channel or device from params when trying to resend
     def service_from_resend_params