update test_dropout

dame-cell · Nov 1, 2024 · 4818d73 · 4818d73
1 parent 9fe754c
commit 4818d73
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -70,6 +70,7 @@ pip install -U triformer
 ```bash
 pytest tests/test_layernorm.py
 pytest tests/test_softmax.py
+pytest tests/test_dropout.py
 ```
 
 ## Future Plans - To Do

diff --git a/tests/test_dropout.py b/tests/test_dropout.py
@@ -1,4 +1,6 @@
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 import pytest
 from triformer import TritonDropout