delete info lazy init in tests

tests/test_block_exact_match.py

@@ -24,7 +24,7 @@ def test_remote_block_exact_match(atol_forward=1e-5, atol_inference=1e-3):
         assert isinstance(remote_block, RemoteTransformerBlock)
 
         inputs = torch.randn(1, 8, config.hidden_size)
-        (outputs_forward,) = remote_block(inputs)
+        outputs_forward = remote_block(inputs)
 
         outputs_inference = []
         with remote_block.inference_session(max_length=inputs.shape[1]) as sess:
@@ -39,7 +39,7 @@ def test_remote_block_exact_match(atol_forward=1e-5, atol_inference=1e-3):
         outputs_inference = torch.cat(outputs_inference, dim=1)
 
         ref_block = load_pretrained_block(MODEL_NAME, block_index, torch_dtype=torch.float32)
-        (outputs_local,) = ref_block(inputs)
+        outputs_local = ref_block(inputs)
 
         assert torch.allclose(outputs_local, outputs_forward, rtol=0, atol=atol_forward)
         assert torch.allclose(outputs_local, outputs_inference, rtol=0, atol=atol_inference)

tests/test_chained_calls.py

@@ -26,9 +26,6 @@ def test_forward_backward_exact_match(atol_forward=1e-4, atol_backward=1e-4, seq
     assert remote_block is not None, f"Could not find {MODEL_NAME}{UID_DELIMITER}0 in DHT"
     assert isinstance(remote_block, RemoteTransformerBlock)
 
-    _ = remote_block.info  # lazy-init info now, because otherwise we will _break_ info init by chaning _info
-    remote_block._info = ExpertInfo(f"{MODEL_NAME}.3 {MODEL_NAME}.4 {MODEL_NAME}.5", remote_block._info.peer_id)
-
     ref_blocks = [
         load_pretrained_block(MODEL_NAME, 3, torch_dtype=torch.float32),
         load_pretrained_block(MODEL_NAME, 4, torch_dtype=torch.float32),
@@ -59,9 +56,6 @@ def test_chained_inference_exact_match(atol_inference=1e-4):
     assert remote_block is not None, f"Could not find {MODEL_NAME}{UID_DELIMITER}0 in DHT"
     assert isinstance(remote_block, RemoteTransformerBlock)
 
-    _ = remote_block.info  # lazy-init info now, because otherwise we will _break_ info init by chaning _info
-    remote_block._info = ExpertInfo(f"{MODEL_NAME}.3 {MODEL_NAME}.4", remote_block._info.peer_id)
-
     inputs = torch.randn(1, 8, config.hidden_size)
 
     outputs_inference = []