code: def partition_dataset():
      """Partitioning MNIST"""
      dataset = datasets.MNIST(
          "./data",
          train=True,
          download=True,
          transform=transforms.Compose(
              [
                  transforms.ToTensor(),
                  transforms.Normalize((0.1307,), (0.3081,)),
              ]
          ),
      )
      size = dist.get_world_size()
      bsz = int(128 / float(size))
      partition_sizes = [1.0 / size for _ in range(size)]
      partition = DataPartitioner(dataset, partition_sizes)
      partition = partition.use(dist.get_rank())
      train_set = torch.utils.data.DataLoader(
          partition, batch_size=bsz, shuffle=True
      )
      return train_set, bsz
  
imports: import math
  import random
  import time
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import torchvision.datasets
  import torchvision.transforms
  import fiber
  from fiber.experimental import Ring
label_desc: Use image node_id 1 for calling a global function with example usage: partition_dataset() and returns: train_set,  bsz