flag to remove optimizer

scripts/run_server.py

@@ -23,6 +23,7 @@ if __name__ == "__main__":
     parser.add_argument('--port', type=int, default=None, required=False)
     parser.add_argument('--host', type=str, default='0.0.0.0', required=False)
     parser.add_argument('--no_network', action='store_true')
+    parser.add_argument('--no_optimizer', action='store_true')
     parser.add_argument('--initial_peers', type=str, default="[]", required=False)
     parser.add_argument('--network_port', type=int, default=None, required=False)
     parser.add_argument('--lifetime_seconds', type=int, default=None, required=False)
@@ -55,9 +56,9 @@ if __name__ == "__main__":
     experts = {}
     for i in range(args.num_experts):
         expert = torch.jit.script(layers.name_to_block[args.expert_cls](args.hidden_dim))
+        opt = torch.optim.SGD(expert.parameters(), 0.0) if args.no_optimizer else torch.optim.Adam(expert.parameters())
         expert_uid = f'{args.expert_prefix}.{i + args.expert_offset}'
-        experts[expert_uid] = tesseract.ExpertBackend(name=expert_uid,
-                                                      expert=expert, opt=torch.optim.Adam(expert.parameters()),
+        experts[expert_uid] = tesseract.ExpertBackend(name=expert_uid, expert=expert, opt=opt,
                                                       args_schema=(tesseract.BatchTensorProto(args.hidden_dim),),
                                                       outputs_schema=tesseract.BatchTensorProto(args.hidden_dim),
                                                       max_batch_size=args.max_batch_size,