I have followed your suggestion and run the gpu code separately in a separate python interpreter as follows. But the effect is same.
import ....
.....
.....
def train(model, device, inputs, labels, epochs):
    print(model.state_dict())
    model.to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
    print(model.state_dict())
    for epoch in range(epochs):
        model.train()
        print(f'Epoch {epoch}')
        optimizer.zero_grad()
        outputs = model(inputs.to(device))
        loss = criterion(outputs, labels.to(device))
        loss.backward()
        optimizer.step()
    print(model.state_dict())
    return loss


X, y = generate_1d_data()
inputs = torch.tensor(X.reshape(-1, 1), dtype=torch.float32)
labels = torch.tensor(y.reshape(-1, 1), dtype=torch.float32)
criterion = nn.MSELoss()

# Running with GPU
torch.cuda.empty_cache()
model = LinearRegressionModel(input_dim=1, output_dim=1)
device_gpu = get_device()

loss = train(model, device_gpu, inputs, labels, epochs=1000)
predictions_gpu = get_prediction(model, device_gpu, inputs)


sns.scatterplot(x=X, y=y, color='blue', label='Data')
sns.lineplot(x=X, y=predictions_gpu.ravel(), color='red', label='Linear Model')
plt.savefig('temp/training.png')

But the prediction always produce constant result from gpu, not sure why.
predictions_gpu
array([[8.821195],
       [8.821195],
       [8.821195],
       [8.821195],
       [8.821195],
       [8.821195],
       [8.821195],
...

But I have checked the parameter state has been changed from cpu to gpu. Initial state OrderedDict([('linear.weight', tensor([[0.0323]])), ('linear.bias', tensor([0.9064]))]) to moving to gpu
OrderedDict([('linear.weight', tensor([[0.0323]], device='cuda:0')), ('linear.bias', tensor([0.9064], device='cuda:0'))]) then final state OrderedDict([('linear.weight', tensor([[145.3438]], device='cuda:0')), ('linear.bias', tensor([8.8212], device='cuda:0'))])
Though I have given the correct input:
inputs
tensor([[ 0.0000],
        [ 0.1005],
        [ 0.2010],
        [ 0.3015],
        [ 0.4020],
        [ 0.5025],
        [ 0.6030],
        [ 0.7035],
        [ 0.8040],
        [ 0.9045],
        [ 1.0050],
.....