测试a100 torch 配合cuda 能否正常运行

测试程序

# -*- coding: utf-8 -*-
def gpu_test():
    """
    python -c "import uutils; uutils.torch_uu.gpu_test()"
    """
    from torch import Tensor
    import torch

    x = torch.randn(2, 4).cuda()
    y = torch.randn(4, 1).cuda()
    out = (x @ y)
    print(out.shape)
    print('Success, no Cuda errors means it worked see:\n')

gpu_test()