Update models/unet.py

models/unet.py

@@ -25,7 +25,7 @@ class CustomLayerNorm(nn.LayerNorm):
 def replace_layer_norm(model):
     for name, module in model.named_children():
         if isinstance(module, nn.LayerNorm):
-            setattr(model, name, CustomLayerNorm(module.normalized_shape, elementwise_affine=module.elementwise_affine))
+            setattr(model, name, CustomLayerNorm(module.normalized_shape, elementwise_affine=module.elementwise_affine)).cuda()
         else:
             replace_layer_norm(module)  # Recursively apply to all submodules