瀏覽代碼

black-isort

justheuristic 2 年之前
父節點
當前提交
c096ebdbbc
共有 1 個文件被更改,包括 3 次插入1 次删除
  1. 3 1
      tests/test_remote_sequential.py

+ 3 - 1
tests/test_remote_sequential.py

@@ -79,7 +79,9 @@ def test_remote_sequential_prompts(batch_size=2, seq_len=5, pre_seq_len=3):
 
         block = load_pretrained_block(MODEL_NAME, block_index=block_index, torch_dtype=torch.float32)
         (outputs_ref,) = block(outputs_ref)
-    outputs_ref = (outputs_ref - torch.cat([inputs, input_prompts_ref], dim=1)) + torch.cat([inputs, input_prompts_ref], dim=1)
+    outputs_ref = (outputs_ref - torch.cat([inputs, input_prompts_ref], dim=1)) + torch.cat(
+        [inputs, input_prompts_ref], dim=1
+    )
 
     assert torch.allclose(outputs_ref, outputs)  # exact match