numpy.array与torch.tensor转换

import numpy as np
from torch import nn
import torch

"""test_on_cpu"""
a = torch.ones(5)
print(a)
print(a.size())
b = a.numpy()
print(b)
print(b.size)
c = torch.from_numpy(b)
print(c)
print(c.size())

tensor([1., 1., 1., 1., 1.])
torch.Size([5])
[1. 1. 1. 1. 1.]
5
tensor([1., 1., 1., 1., 1.])
torch.Size([5])

"""test_on_gpu"""
如果a从gpu获得
print(a.size())
b = a.data.cpu().numpy()
print(b.size)

torch.Size([1, 9216])
9216