Spaces:
Running on Zero
Running on Zero
| # Copyright (c) 2021 OpenAI | |
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: MIT AND Apache-2.0 | |
| # This code is derived from https://github.com/openai/guided-diffusion | |
| """ | |
| Various utilities for neural networks. | |
| """ | |
| def mean_flat(tensor): | |
| """ | |
| Take the mean over all non-batch dimensions. | |
| """ | |
| return tensor.mean(dim=list(range(1, len(tensor.shape)))) | |
| def sum_flat(tensor): | |
| """ | |
| Take the sum over all non-batch dimensions. | |
| """ | |
| return tensor.sum(dim=list(range(1, len(tensor.shape)))) | |