18 lines
498 B
Python
18 lines
498 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class PadCrop(nn.Module):
|
|
def __init__(self, crop_size=None, pad_mode='constant', pad_value=0):
|
|
super().__init__()
|
|
self.crop_size = crop_size
|
|
self.pad_mode = pad_mode
|
|
self.pad_value = pad_value
|
|
|
|
def forward(self, x):
|
|
# Simple implementation - you can customize this
|
|
if self.crop_size:
|
|
return x[:, :, :self.crop_size]
|
|
return x
|
|
|
|
# Export PadCrop
|
|
__all__ = ['PadCrop'] |