From 2ecd61e3e2e322acf10135aed6142d2651cbccab Mon Sep 17 00:00:00 2001
From: Elias Wahl <82230675+Eliulm@users.noreply.github.com>
Date: Fri, 19 Apr 2024 01:20:52 +0200
Subject: [PATCH] monkey patching (#4214)

---
 extra/models/bert.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/extra/models/bert.py b/extra/models/bert.py
index 787a2c261b..0b95773454 100644
--- a/extra/models/bert.py
+++ b/extra/models/bert.py
@@ -1,8 +1,12 @@
 from tinygrad.tensor import Tensor
-from tinygrad.nn import Linear, LayerNorm, Embedding
+from tinygrad import nn
 from tinygrad.helpers import fetch, get_child
 from pathlib import Path
 
+# allow for monkeypatching
+Embedding = nn.Embedding
+Linear = nn.Linear
+LayerNorm = nn.LayerNorm
 
 class BertForQuestionAnswering:
   def __init__(self, hidden_size=1024, intermediate_size=4096, max_position_embeddings=512, num_attention_heads=16, num_hidden_layers=24, type_vocab_size=2, vocab_size=30522, attention_probs_dropout_prob=0.1, hidden_dropout_prob=0.1):