test_gpu.py
import tntorch as tn
import torch
torch.set_default_dtype(torch.float64)
# in case the computer testing has no gpu, tests will just pass
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def test_tt():
X = torch.randn(16, 16, 16)
y1 = tn.Tensor(X, ranks_tt=3).torch()
y2 = tn.Tensor(X, ranks_tt=3, device=device).torch().cpu()
assert torch.abs(y1-y2).max() < 1e-5
def test_tucker():
X = torch.randn(16, 16, 16)
y1 = tn.Tensor(X, ranks_tucker=3).torch()
y2 = tn.Tensor(X, ranks_tucker=3, device=device).torch().cpu()
assert torch.abs(y1-y2).max() < 1e-5
def test_cp():
X = torch.randn(16, 16, 16)
y1 = tn.Tensor(X, ranks_cp=3).torch()
y2 = tn.Tensor(X, ranks_cp=3, device=device).torch().cpu()
assert torch.abs(y1-y2).max() < 1e-5