Add some group parameter to convolutions. (#566)

* Add some group parameter to convolutions.

* Avoid some unnecessary groups checks.

* Move the tensor convolution bits.

* Properh handling of groups.

* Bump the crate version.

* And add a changelog.
This commit is contained in:
Laurent Mazare
2023-08-23 12:58:55 +01:00
committed by GitHub
parent 4ee1cf038a
commit aba1e90797
30 changed files with 216 additions and 113 deletions

View File

@ -112,8 +112,8 @@ impl UNet2DConditionModel {
let bl_attention_head_dim = config.blocks.last().unwrap().attention_head_dim;
let time_embed_dim = b_channels * 4;
let conv_cfg = nn::Conv2dConfig {
stride: 1,
padding: 1,
..Default::default()
};
let conv_in = conv2d(in_channels, b_channels, 3, conv_cfg, vs.pp("conv_in"))?;