PyTorch create tensors directly on device Medium

Creating PyTorch tensors on the CPU and then moving them to the device impacts the performance.

Detector ID
python/pytorch-create-tensors-directly-on-device@v1.0
Category
Common Weakness Enumeration (CWE) external icon
-

Noncompliant example

1def pytorch_create_tensors_directly_on_device_noncompliant():
2    import torch
3    t = torch.ones(list(range(1, 11)), dtype=torch.float64)
4    # Noncompliant: tensor is created on cpu and then moved to device.
5    t.cuda()

Compliant example

1def pytorch_create_tensors_directly_on_device_compliant():
2    import torch
3    # Compliant: tensor is directly created on device.
4    t = torch.tensor([1, 2, 3, 4], device="cuda")