Why am I not able to use directml with PyTorch?

Question

I am trying to use DirectML with PyTorch to train a neural network as my PC doesn't support CUDA.

I am a newbie getting started with PyTorch. I am trying to run the following code using DirectML.

import torch
import torch_directml
from torch.utils.tensorboard import SummaryWriter

dml = torch_directml.device(torch_directml.default_device())
writer = SummaryWriter()

x = torch.arange(-5, 5, 0.1).view(-1, 1)
y = -5 * x + 0.1 * torch.randn(x.size())
x = x.to(dml)
y = y.to(dml)

model = torch.nn.Linear(1, 1).to(dml)
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr = 0.1)

def train_model(iter):
    for epoch in range(iter):
        y1 = model(x)
        loss = criterion(y1, y)
        writer.add_scalar("Loss/train", loss, epoch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

train_model(10)
writer.flush()

I am getting the following error:
RuntimeError: 0 <= device.index() && device.index() < static_cast<c10::DeviceIndex>(device_ready_queues_.size()) INTERNAL ASSERT FAILED at "..\\torch\\csrc\\autograd\\engine.cpp":1372, please report a bug to PyTorch.

Can anyone help me with this. Thanks in advance.