diff --git a/S1/9/rmsnorm_cuda.py b/S1/9/rmsnorm_cuda.py
index 904f3cc..8a0face 100644
--- a/S1/9/rmsnorm_cuda.py
+++ b/S1/9/rmsnorm_cuda.py
@@ -1,4 +1,5 @@
 import torch
+import torch.nn as nn
 from torch.utils.cpp_extension import load_inline
 
 rmsnorm_source = """