torchscript_consistency_cpu_test.py 634 Bytes
Newer Older
1
import torch
2

3
from . import common_utils
4
from .torchscript_consistency_impl import Functional, Transforms
5
6


moto's avatar
moto committed
7
class TestFunctionalFloat32(Functional, common_utils.PytorchTestCase):
8
9
    dtype = torch.float32
    device = torch.device('cpu')
10
11


moto's avatar
moto committed
12
class TestFunctionalFloat64(Functional, common_utils.PytorchTestCase):
13
14
    dtype = torch.float64
    device = torch.device('cpu')
15
16


moto's avatar
moto committed
17
class TestTransformsFloat32(Transforms, common_utils.PytorchTestCase):
18
19
20
21
    dtype = torch.float32
    device = torch.device('cpu')


moto's avatar
moto committed
22
class TestTransformsFloat64(Transforms, common_utils.PytorchTestCase):
23
24
    dtype = torch.float64
    device = torch.device('cpu')