change some stuff

cudagraph.py

@@ -65,8 +65,8 @@ def test_thing(graph, input):
     torch.cuda.synchronize()
 
 with torch.no_grad():
-    model = init_1_3b().cuda().half()
-    shape = (1, 256)
+    model = init_6b().cuda().half()
+    shape = (1, 1)
     x = torch.zeros(shape).cuda().long()
     print(shape)
     print("PyTorch Eager")

main.py

@@ -225,12 +225,13 @@ class GPTLayer(nn.Module):
     def forward(self, x, hypernetwork):
         residual = x
         x = self.ln_preattn(x)
-        if hypernetwork:
-            hyper_out = hypernetwork(x)
-
         attn_out = self.attn(x)
         ff_out = self.ff(x)
-        x = residual + ff_out + attn_out# + (hyper_out if hypernetwork is not None else 0)
+        x = residual + ff_out + attn_out
+        if hypernetwork:
+            hyper_out = hypernetwork(x)
+            x = x + hyper_out
+            
         return x
 
 # Can access and change every module from here, as both Layer class and ff and attn classes are passed from GPTModel.

test_pyfra.py

@@ -18,9 +18,6 @@ config_obj.create_service(overwrite=True)
 remote = config_obj.get_pyfra_remote()
 env1 = remote.env('noname', python_version=None)
 
-models = {'6b': '/home/xuser/models/j6b_ckpt_14001', '20b': '/home/xuser/diffusionstorage/models/20B_136500.hf'
-          ,'13b': '/home/xuser/diffusionstorage/workspace/finetune/pokepls/bigmodel'}
-
 path = env1.path('/home/xuser/diffusionstorage/workspace/kuru/basedformer')
 env1.sh('pip install /home/xuser/hugessd/pytorch/torch-1.10.1+cu113-cp38-cp38-linux_x86_64.whl')
 env1.sh('pip install einops numpy')