torchscript_consistency_cpu_test.py 606 Bytes
Newer Older
1
import torch
2

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


7
8
9
class TestFunctionalFloat32(Functional, common_utils.TestCase):
    dtype = torch.float32
    device = torch.device('cpu')
10
11


12
13
14
class TestFunctionalFloat64(Functional, common_utils.TestCase):
    dtype = torch.float64
    device = torch.device('cpu')
15
16


17
18
19
20
21
22
23
24
class TestTransformsFloat32(Transforms, common_utils.TestCase):
    dtype = torch.float32
    device = torch.device('cpu')


class TestTransformsFloat64(Transforms, common_utils.TestCase):
    dtype = torch.float64
    device = torch.device('cpu')