Update seed2_tokenizer.py

seed2_tokenizer.py

@@ -1654,7 +1654,7 @@ class Blip2Base(PreTrainedModel):
 
     @classmethod
     def init_Qformer(cls, encoder_config, num_query_token, vision_width, cross_attention_freq=2, cache_dir=""):
-        #print ("loading")
+        print ("loading")
         encoder_config = BertConfig.from_pretrained("bert-base-uncased")
         encoder_config.encoder_width = vision_width
         # insert cross-attention layer every other block
@@ -2091,10 +2091,11 @@ class Seed2Tokenizer(PreTrainedModel):
         ])
 
         shape_latents = torch.Size([1, 4, 96, 96])
-        self.latents = torch.randn(shape_latents, generator=None, layout=torch.strided)
+        self.register_buffer("latents",torch.randn(shape_latents, generator=None, layout=torch.strided))
+        
 
         shape_noise = torch.Size([1, 1024])
-        self.noise = torch.randn(shape_noise, generator=None, layout=torch.strided)
+        self.register_buffer("noise",torch.randn(shape_noise, generator=None, layout=torch.strided))
 
         self.model = model
         self.processor = processor
@@ -2183,5 +2184,4 @@ class Seed2Tokenizer(PreTrainedModel):
 
             image_torch = image_torch.to(self.device)
         return self.encode(image_torch, visual_encoder)
-    
-
+