comparing_graphs.py

import torch
import torch.nn as nn
from torchviz import make_dot

device = 'cuda'
# sending tensor to device at creation time
a = torch.randn(1, requires_grad=True, dtype=torch.float, device=device)
plot1 = make_dot(a)

# sending tensor to device immediately after creating it
a = torch.randn(1, requires_grad=True, dtype=torch.float).to(device)
plot2 = make_dot(a)

class ManualLinearRegression(nn.Module):
    def __init__(self):
        super().__init__()
        # To make "a" and "b" real parameters of the model, we need to wrap them with nn.Parameter
        self.a = nn.Parameter(torch.randn(1, requires_grad=True, dtype=torch.float))
        self.b = nn.Parameter(torch.randn(1, requires_grad=True, dtype=torch.float))
        
    def forward(self, x):
        # Computes the outputs / predictions
        return self.a + self.b * x
    
# sending model to device
model = ManualLinearRegression().to(device)
plot3 = make_dot(model.a)

# plots n.1 and n.3 will be exactly the same
# plot n.2, on the other hand, is different - the `to` operation created an unwanted computation graph