[Type Hint] VAE models (#365)

daspartho · anton-l · natolambert · commit c37a5ed08e2f · 2022-09-07T09:35:06.000-06:00
* [Type Hint] VAE models

* Update src/diffusers/models/vae.py

* apply suggestions from code review

Co-authored-by: Anton Lozhkov &lt;aglozhkov@gmail.com&gt;
diff --git a/src/diffusers/models/vae.py b/src/diffusers/models/vae.py
@@ -411,7 +411,7 @@ def __init__(
             act_fn=act_fn,
         )
 
-    def encode(self, x, return_dict: bool = True):
+    def encode(self, x: torch.FloatTensor, return_dict: bool = True) -> VQEncoderOutput:
         h = self.encoder(x)
         h = self.quant_conv(h)
 

Original file line number	Diff line number	Diff line change
`@@ -411,7 +411,7 @@ def __init__(`
`411`	`411`	`act_fn=act_fn,`
`412`	`412`	`)`
`413`	`413`
`414`		`- def encode(self, x, return_dict: bool = True):`
	`414`	`+ def encode(self, x: torch.FloatTensor, return_dict: bool = True) -> VQEncoderOutput:`
`415`	`415`	`h = self.encoder(x)`
`416`	`416`	`h = self.quant_conv(h)`
`417`	`417`