Skip to content

GitLab

Unverified Commit 0c0c2224 authored by Mishig Davaadorj's avatar Mishig Davaadorj Committed by GitHub
Browse files

FlaxUNet2DConditionOutput @flax.struct.dataclass (#550)

parent d09bbae5
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
src/diffusers/models/unet_2d_condition_flax.py
View file @ 0c0c2224
from dataclasses import dataclass
from typing import Tuple, Union
import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
......@@ -19,7 +19,7 @@ from .unet_blocks_flax import (
)
@dataclass
@flax.struct.dataclass
class FlaxUNet2DConditionOutput(BaseOutput):
"""
Args:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment