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
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")