update main.rs

ferritin-examples/examples/amplify/main.rs

@@ -88,16 +88,20 @@ fn main() -> Result<()> {
 
         let token_ids = Tensor::new(&tokens[..], device)?.unsqueeze(0)?;
         println!("Encoding.......");
-        let encoded = model.forward(&token_ids, None, false, false)?;
+        let encoded = model.forward(&token_ids, None, false, true)?;
 
         println!("Predicting.......");
         let predictions = encoded.logits.argmax(D::Minus1)?;
+        println!("Pred Dims: {:?}", encoded.logits.dims());
 
         println!("Decoding.......");
         let indices: Vec<u32> = predictions.to_vec2()?[0].to_vec();
         let decoded = tokenizer.decode(indices.as_slice(), true);
 
         println!("Decoded: {:?}, ", decoded);
+
+        let contact_map = encoded.get_contact_map()?;
+        println!("Contact Map Calculated: {:?}, ", contact_map);
     }
 
     Ok(())