Use rotary positional embeddings in relformer_imagenet64.gin config

wolf-lol · Aug 15, 2021 · 753d73f · 753d73f
1 parent 5706960
commit 753d73f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/trax/supervised/configs/relformer_imagenet64.gin b/trax/supervised/configs/relformer_imagenet64.gin
@@ -66,9 +66,10 @@ train.checkpoints_at = \
 
 # Parameters for PureLSHSelfAttentionWrapper:
 # ==============================================================================
+PureLSHSelfAttentionWrapper.pure_lsh_implementation = @PureLSHSelfAttention
+PureLSHSelfAttentionWrapper.rotary_position_emb = True
 PureLSHSelfAttentionWrapper.bias = True
 PureLSHSelfAttentionWrapper.num_weights = 2
-PureLSHSelfAttentionWrapper.pure_lsh_implementation = @PureLSHSelfAttention
 PureLSHSelfAttentionWrapper.weights_format = 'model'
 
 # Parameters for PureLSHSelfAttention: