From 19e41595cc6f932acadd5eed982873e3423762bb Mon Sep 17 00:00:00 2001
From: tuner007 <46425391+tuner007@users.noreply.github.com>
Date: Thu, 17 Jun 2021 02:48:49 +0530
Subject: [PATCH] Update models.py

Missing 1 required positional argument "device"
---
 models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models.py b/models.py
index cb2e90f..9616c8d 100644
--- a/models.py
+++ b/models.py
@@ -57,7 +57,7 @@ class CausalLinearAttentionAMP(CausalLinearAttention):
 
     def forward(self, queries, keys, values, query_mask=None,
                 key_mask=None, cache=None):
-        self.feature_map.new_feature_map()
+        self.feature_map.new_feature_map(queries.device)
         Q = self.feature_map.forward_queries(queries)
         K = self.feature_map.forward_keys(keys)