trw.layers.flatten

Module Contents

Classes

Flatten

Flatten a tensor

class trw.layers.flatten.Flatten

Bases: torch.nn.Module

Flatten a tensor

For example, a tensor of shape[N, Z, Y, X] will be reshaped [N, Z * Y * X]

forward(self, x: torch.Tensor) → torch.Tensor

Parameters

x – a tensor

Returns: return a flattened tensor