Certainly. Sorry for missing out the details. Here is my full code. One strange thing is I was executing my code now to make sure I provided the right one and it worked without issues. I am certain the same code threw the error message with CUDA devices the last time I ran.
Below classes is for my model. The Flow() class creates the entire model.
import torch 
from torch import nn

from sys import exit as e

class SimpleNet(nn.Module):
  def __init__(self, inp, parity):
    super(SimpleNet, self).__init__()
    self.net = nn.Sequential(
      nn.Linear(inp//2, 256),
      nn.LeakyReLU(True),
      nn.Linear(256, 256),
      nn.LeakyReLU(True),
      nn.Linear(256, inp//2),
      nn.Sigmoid()
    )
    self.inp = inp
    self.parity = parity
  
  def forward(self, x):
    z = torch.zeros_like(x)
    x0, x1 = x[:, :, ::2, ::2], x[:, :, 1::2, 1::2]
    if self.parity % 2:
      x0, x1 = x1, x0 
    z1 = x1
    log_s = self.net(x1)
    t = self.net(x1)
    s = torch.exp(log_s)
    z0 = (s * x0) + t
    if self.parity%2:
      z0, z1 = z1, z0
    z[:, :, ::2, ::2] = z0
    z[:, :, 1::2, 1::2] = z1
    logdet = torch.sum(torch.log(s), dim = 1)
    return z, logdet
  
  def reverse(self, z):
    x = torch.zeros_like(z)
    z0, z1 = z[:, :, ::2, ::2], z[:, :, 1::2, 1::2]
    if self.parity%2:
      z0, z1 = z1, z0
    x1 = z1
    log_s = self.net(z1)
    t = self.net(z1)
    s = torch.exp(log_s)
    x0 = (z0 - t)/s
    if self.parity%2:
      x0, x1 = x1, x0
    x[:, :, ::2, ::2] = x0
    x[:, :, 1::2, 1::2] = x1
    return x
  

class Block(nn.Module):
  def __init__(self, inp, n_blocks):
    super(Block, self).__init__()
    parity = 0
    self.blocks = nn.ModuleList()
    for _ in range(n_blocks):
      self.blocks.append(SimpleNet(inp, parity))
      parity += 1
    
  
  def forward(self, x):
    logdet = 0
    out = x
    xs = [out]
    for block in self.blocks:
      out, det = block(out)
      logdet += det
      xs.append(out)
    return out, logdet

  def reverse(self, z):
    out = z
    for block in self.blocks[::-1]:
      out = block.reverse(out)
    return out


class Flow(nn.Module):
  def __init__(self, inp, prior, n_blocks):
    super(Flow, self).__init__()
    self.prior = prior
    self.flow = Block(inp, n_blocks)
  
  def forward(self, x):
    z, logdet = self.flow(x)
    logprob = self.prior.log_prob(z).view(x.size(0), -1).sum(1) #Error encountered here
    return z, logdet, logprob
  
  def reverse(self, z):
    x = self.flow.reverse(z)
    return x
  
  def get_sample(self, n):
    z = self.prior.sample(sample_shape = torch.Size([n]))
    return self.reverse(z)


I define my prior variable and instantiate the model as follows and train it on MNIST dataset
def startup(opt, use_cuda):
  torch.manual_seed(1)
  device = "cuda" if not opt.no_cuda and torch.cuda.is_available() else "cpu"
  kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}

  inpt = 100
  dim = 1
  img_size = 28
  n_block = 4
  epochs = 50
  lr = 0.01
  wd=1e-3
  old_loss = 1e6
  best_loss = 0
  batch_size = 128
  prior = MultivariateNormal(torch.zeros(img_size).to(device), torch.eye(img_size).to(device))

  #MNIST
  transform = transforms.Compose([transforms.ToTensor()])
  train_dataset = MNIST(root=opt.root, train=True, transform=transform, \
    download=True)
  train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, **kwargs)
  
  model = Flow(dim, prior, n_block)
  model = model.to(device)

  if use_cuda and torch.cuda.device_count()>1:
        model = DataParallel(model, device_ids=[0, 1, 2, 3])

  optimizer = optim.Adam(model.parameters(), lr)
  scheduler = StepLR(optimizer, step_size=1, gamma=0.7)
  
  t0 = time.time()
  for epoch in range(epochs):
    model.train()
    train_data(opt, model, device, train_loader, optimizer, epoch)
    scheduler.step()
  print(f"time to complete {epochs} epoch: {time.time() - t0} seconds")


Training is done the usual way with some preprocessing
def preprocess(x):
  x = x * 255
  x = torch.floor(x/2**3)
  x = x/32 - 0.5
  return x

for b, (x, _) in enumerate(train_loader):
    optimizer.zero_grad()
    x = x.to(device)
    x = preprocess(x)
    x = x.view(x.size(0), -1)
    z, logdet, logprob = model.module(x) # Encounters the original error

(I have skipped the rest of the steps as the program stops here)