Преглед на файлове

feat: modify main for fedssl

shellmiao преди 1 година
родител
ревизия
b7f7b217ba
променени са 1 файла, в които са добавени 4 реда и са изтрити 2 реда
  1. 4 2
      applications/fedssl/main.py

+ 4 - 2
applications/fedssl/main.py

@@ -1,4 +1,6 @@
 import argparse
+from applications.fedssl.client_with_pgfed import FedSSLWithPgFedClient
+from applications.fedssl.server_with_pgfed import FedSSLWithPgFedServer
 
 import easyfl
 from client import FedSSLClient
@@ -162,8 +164,8 @@ def run():
 
     model = get_model(args.model, args.encoder_network, args.predictor_network)
     easyfl.register_model(model)
-    easyfl.register_client(FedSSLClient)
-    easyfl.register_server(FedSSLServer)
+    easyfl.register_client(FedSSLWithPgFedClient)
+    easyfl.register_server(FedSSLWithPgFedServer)
     easyfl.init(config, init_all=True)
     easyfl.run()