use omni endpoint

lightly-ai · Jan 22, 2024 · d12c59f · d12c59f
1 parent 17e1038
commit d12c59f
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/lightly/api/api_workflow_compute_worker.py b/lightly/api/api_workflow_compute_worker.py
@@ -17,7 +17,7 @@
     DockerRunScheduledState,
     DockerRunState,
     DockerWorkerConfigV3,
-    DockerWorkerConfigV3CreateRequest,
+    DockerWorkerConfigOmniVXCreateRequest,
     DockerWorkerConfigV3Docker,
     DockerWorkerConfigV3Lightly,
     DockerWorkerRegistryEntryData,
@@ -243,11 +243,13 @@ def create_compute_worker_config(
             lightly=lightly,
             selection=selection,
         )
-        request = DockerWorkerConfigV3CreateRequest(
-            config=config, creator=self._creator
-        )
+        request = DockerWorkerConfigOmniVXCreateRequest.from_dict({
+            "version": "V3",
+            "config": config.to_dict(by_alias=True),
+            "creator": self._creator,
+        })
         try:
-            response = self._compute_worker_api.create_docker_worker_config_v3(request)
+            response = self._compute_worker_api.create_docker_worker_config_vx(request)
             return response.id
         except ApiException as e:
             if e.body is None: