Fix format

Author: lantiga

examples/fabric/tensor_parallel/train.py

@@ -1,13 +1,14 @@
 import lightning as L
 import torch
 import torch.nn.functional as F
-from data import RandomTokenDataset
 from lightning.fabric.strategies import ModelParallelStrategy
 from model import ModelArgs, Transformer
 from parallelism import parallelize
 from torch.distributed.tensor.parallel import loss_parallel
 from torch.utils.data import DataLoader
 
+from data import RandomTokenDataset
+
 
 def train():
     strategy = ModelParallelStrategy(

examples/pytorch/tensor_parallel/train.py

@@ -1,13 +1,14 @@
 import lightning as L
 import torch
 import torch.nn.functional as F
-from data import RandomTokenDataset
 from lightning.pytorch.strategies import ModelParallelStrategy
 from model import ModelArgs, Transformer
 from parallelism import parallelize
 from torch.distributed.tensor.parallel import loss_parallel
 from torch.utils.data import DataLoader
 
+from data import RandomTokenDataset
+
 
 class Llama3(L.LightningModule):
     def __init__(self):