daviddoji/Made-With-ML
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

updated to Ray 2.7

Browse Source
This commit is contained in:
GokuMohandas
2023-09-18 22:03:20 -07:00
parent 71b3d50a05
commit b98bd5b1ae
15 changed files with 3484 additions and 2086 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/code/test_utils.py
View File

@@ -4,6 +4,7 @@ from pathlib import Path
import numpy as np
import pytest
import torch
from ray.train.torch import get_device
from madewithml import utils
@@ -42,9 +43,9 @@ def test_collate_fn():
}
processed_batch = utils.collate_fn(batch)
expected_batch = {
"ids": torch.tensor([[1, 2, 0], [1, 2, 3]], dtype=torch.int32),
"masks": torch.tensor([[1, 1, 0], [1, 1, 1]], dtype=torch.int32),
"targets": torch.tensor([3, 1], dtype=torch.int64),
"ids": torch.as_tensor([[1, 2, 0], [1, 2, 3]], dtype=torch.int32, device=get_device()),
"masks": torch.as_tensor([[1, 1, 0], [1, 1, 1]], dtype=torch.int32, device=get_device()),
"targets": torch.as_tensor([3, 1], dtype=torch.int64, device=get_device()),
}
for k in batch:
assert torch.allclose(processed_batch[k], expected_batch[k])