From 2042319e355fc5ec11a40d62f1079a5f5a752d99 Mon Sep 17 00:00:00 2001
From: eclipsevortex <eclipsevortext@hotmail.com>
Date: Tue, 3 Dec 2024 18:49:24 +0000
Subject: [PATCH] fix set weight and websocket connection

---
 neurons/validator.py          |  1 +
 subnet/validator/challenge.py | 13 ++++---------
 subnet/validator/weights.py   |  3 ++-
 3 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/neurons/validator.py b/neurons/validator.py
index 44bd1d55..a71e8317 100644
--- a/neurons/validator.py
+++ b/neurons/validator.py
@@ -297,6 +297,7 @@ async def run(self):
                     btul.logging.debug(f"Setting weights {self.moving_averaged_scores}")
                     set_weights_for_validator(
                         uid=self.uid,
+                        device=self.device,
                         subtensor=self.subtensor,
                         wallet=self.wallet,
                         metagraph=self.metagraph,
diff --git a/subnet/validator/challenge.py b/subnet/validator/challenge.py
index fa55f803..5227d8a6 100644
--- a/subnet/validator/challenge.py
+++ b/subnet/validator/challenge.py
@@ -23,7 +23,6 @@
 import bittensor.core.settings as btcse
 import bittensor.utils.btlogging as btul
 from substrateinterface.base import SubstrateInterface
-from websockets.sync import client as ws_client
 
 from subnet.constants import DEFAULT_PROCESS_TIME
 from subnet.protocol import Synapse
@@ -161,25 +160,21 @@ def challenge_subtensor(miner: Miner, challenge):
 
         # Attempt to connect to the subtensor
         try:
-            # Create the websocket
-            websocket = ws_client.connect(
-                f"ws://{miner.ip}:9944",
-                open_timeout=DEFAULT_PROCESS_TIME,
-                max_size=2**32,
-            )
-
             # Create the substrate
             substrate = SubstrateInterface(
                 ss58_format=btcse.SS58_FORMAT,
                 use_remote_preset=True,
                 type_registry=btcse.TYPE_REGISTRY,
-                websocket=websocket,
+                url=f"ws://{miner.ip}:9944",
             )
 
         except Exception:
             reason = "Failed to connect to Subtensor node at the given IP."
             return (verified, reason, process_time)
 
+        # Set the socket timeout
+        substrate.websocket.settimeout(DEFAULT_PROCESS_TIME)
+
         # Start the timer
         start_time = time.time()
 
diff --git a/subnet/validator/weights.py b/subnet/validator/weights.py
index 55f2c1b3..23e9e2a5 100644
--- a/subnet/validator/weights.py
+++ b/subnet/validator/weights.py
@@ -27,6 +27,7 @@
 
 def set_weights_for_validator(
     uid: int,
+    device: torch.device,
     subtensor: "btcs.Subtensor",
     wallet: "btw.Wallet",
     netuid: int,
@@ -115,7 +116,7 @@ def set_weights_for_validator(
         processed_weight_uids,
         processed_weights,
     ) = btuw.process_weights_for_netuid(
-        uids=metagraph.uids,
+        uids=metagraph.uids.to(device),
         weights=raw_weights,
         netuid=netuid,
         subtensor=subtensor,