Add first successful sampling implementation

2023-08-04 17:01:44 -04:00
parent 3bf5b8c477
commit 77f30aefa7
9 changed files with 42 additions and 34 deletions
--- a/src/model/clip/mod.rs
+++ b/src/model/clip/mod.rs
@@ -59,7 +59,7 @@ impl<B: Backend> CLIP<B> {
    pub fn forward(&self, x: Tensor<B, 2, Int>) -> Tensor<B, 3> {
        let [n_batch, seq_len] = x.dims();
        
-        let mask = attn_decoder_mask(seq_len);
+        let mask = attn_decoder_mask(seq_len, &x.device());

        let embedded = self.token_embedding.forward(x) 
            + self.position_embedding.val().slice([0..seq_len]).unsqueeze();