Update models/DiT.py

Added PyTorchModelHubMixin to allow from_pretrained method

models/DiT.py

@@ -3,6 +3,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import math
 from timm.models.vision_transformer import PatchEmbed
+from huggingface_hub import PyTorchModelHubMixin
 
 class TimestepEmbedder(nn.Module):
     """Module to create timestep's embedding."""
@@ -65,7 +66,8 @@ class DiTBlock(nn.Module):
         x = x * (1+gamma_2.unsqueeze(1)) + beta_2.unsqueeze(1)
         return x
 
-class DiT(nn.Module):
+class DiT(nn.Module,
+         PyTorchModelHubMixin):
     def __init__(self,
                  num_blocks=10,
                  hidden_size=640,