Spaces:
Running on Zero
Running on Zero
File size: 592 Bytes
49d36c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | # 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))))
|