File size: 550 Bytes
e873e70 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | import platform
import torch
import torch.nn.functional as F
import relu
def test_relu():
if platform.system() == "Darwin":
device = torch.device("mps")
elif hasattr(torch, "xpu") and torch.xpu.is_available():
device = torch.device("xpu")
elif torch.version.cuda is not None and torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
x = torch.randn(1024, 1024, dtype=torch.float32, device=device)
torch.testing.assert_allclose(F.relu(x), relu.relu(x))
|