fixing ruff check error

Author: eltsai

src/maxdiffusion/tests/wan_vae_test.py

@@ -17,7 +17,6 @@
 import os
 import functools
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
 import jax
 import jax.numpy as jnp