diff --git a/python/fedml/__init__.py b/python/fedml/__init__.py
index 0e89f8c6a4..2a77abbcfa 100644
--- a/python/fedml/__init__.py
+++ b/python/fedml/__init__.py
@@ -23,7 +23,7 @@
 _global_training_type = None
 _global_comm_backend = None
 
-__version__ = "0.7.340"
+__version__ = "0.7.341"
 
 
 def init(args=None):
diff --git a/python/fedml/cli/edge_deployment/client_diagnosis.py b/python/fedml/cli/edge_deployment/client_diagnosis.py
index f7f76f0370..c42807d7cf 100644
--- a/python/fedml/cli/edge_deployment/client_diagnosis.py
+++ b/python/fedml/cli/edge_deployment/client_diagnosis.py
@@ -68,6 +68,8 @@ def check_mqtt_connection():
                     return False;
                 time.sleep(1)
 
+            mqtt_mgr.disconnect()
+            mqtt_mgr.loop_stop()
             return True
         except Exception as e:
             print("MQTT connect exception: {}".format(str(e)))
diff --git a/python/fedml/core/distributed/communication/mqtt_s3/mqtt_s3_multi_clients_comm_manager.py b/python/fedml/core/distributed/communication/mqtt_s3/mqtt_s3_multi_clients_comm_manager.py
index 694a3bb1b3..d09f9044c0 100755
--- a/python/fedml/core/distributed/communication/mqtt_s3/mqtt_s3_multi_clients_comm_manager.py
+++ b/python/fedml/core/distributed/communication/mqtt_s3/mqtt_s3_multi_clients_comm_manager.py
@@ -198,6 +198,9 @@ def _on_message_impl(self, msg):
                 "mqtt_s3.on_message: model params length %d" % len(model_params)
             )
 
+            model_url = payload_obj.get(Message.MSG_ARG_KEY_MODEL_PARAMS_URL, "")
+            logging.info("mqtt_s3.on_message: model url {}".format(model_url))
+
             # replace the S3 object key with raw model params
             payload_obj[Message.MSG_ARG_KEY_MODEL_PARAMS] = model_params
         else:
diff --git a/python/setup.py b/python/setup.py
index 2eae2cf746..56362b86c5 100755
--- a/python/setup.py
+++ b/python/setup.py
@@ -75,7 +75,7 @@ def finalize_options(self):
 
 setup(
     name="fedml",
-    version="0.7.340",
+    version="0.7.341",
     author="FedML Team",
     author_email="ch@fedml.ai",
     description="A research and production integrated edge-cloud library for "